Buckets:
| <meta charset="utf-8" /><meta name="hf:doc:metadata" content="{"title":"Mask Generation","local":"mask-generation","sections":[{"title":"Mask Generation Pipeline","local":"mask-generation-pipeline","sections":[],"depth":2},{"title":"Model Inference","local":"model-inference","sections":[{"title":"Point Prompting","local":"point-prompting","sections":[],"depth":3},{"title":"Box Prompting","local":"box-prompting","sections":[],"depth":3}],"depth":2},{"title":"Fine-tuning for Mask Generation","local":"fine-tuning-for-mask-generation","sections":[],"depth":2}],"depth":1}"> | |
| <link href="/docs/transformers/pr_39895/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/entry/start.ed44d8d9.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/scheduler.31fdf58d.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/singletons.6a6316d3.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/index.252883d5.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/paths.6325da18.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/entry/app.75a9946c.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/preload-helper.f369a5fb.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/index.2f76fdf0.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/nodes/0.4d11f797.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/each.e59479a4.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/nodes/680.4a15eaef.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/CopyLLMTxtMenu.121ef535.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/MermaidChart.svelte_svelte_type_style_lang.92a7ab29.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/IconCopy.ac192424.js"> | |
| <link rel="modulepreload" href="/docs/transformers/pr_39895/en/_app/immutable/chunks/CodeBlock.e52df5d6.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{"title":"Mask Generation","local":"mask-generation","sections":[{"title":"Mask Generation Pipeline","local":"mask-generation-pipeline","sections":[],"depth":2},{"title":"Model Inference","local":"model-inference","sections":[{"title":"Point Prompting","local":"point-prompting","sections":[],"depth":3},{"title":"Box Prompting","local":"box-prompting","sections":[],"depth":3}],"depth":2},{"title":"Fine-tuning for Mask Generation","local":"fine-tuning-for-mask-generation","sections":[],"depth":2}],"depth":1}"><!-- HEAD_svelte-u9bgzb_END --> <p></p> <div class="items-center shrink-0 min-w-[100px] max-sm:min-w-[50px] justify-end ml-auto flex" style="float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"><div class="inline-flex rounded-md max-sm:rounded-sm"><button class="inline-flex items-center gap-1 h-7 max-sm:h-7 px-2 max-sm:px-1.5 text-sm font-medium text-gray-800 border border-r-0 rounded-l-md max-sm:rounded-l-sm border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-live="polite"><span class="inline-flex items-center justify-center rounded-md p-0.5 max-sm:p-0 hover:text-gray-800 dark:hover:text-gray-200"><svg class="sm:size-3.5 size-3" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg></span> <span>Copy page</span></button> <button class="inline-flex items-center justify-center w-6 max-sm:w-5 h-7 max-sm:h-7 disabled:pointer-events-none text-sm text-gray-500 hover:text-gray-700 dark:hover:text-white rounded-r-md max-sm:rounded-r-sm border border-l transition border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-haspopup="menu" aria-expanded="false" aria-label="Open copy menu"><svg class="transition-transform text-gray-400 overflow-visible sm:size-3.5 size-3 rotate-0" width="1em" height="1em" viewBox="0 0 12 7" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M1 1L6 6L11 1" stroke="currentColor"></path></svg></button></div> </div> <h1 class="relative group"><a id="mask-generation" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#mask-generation"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Mask Generation</span></h1> <p data-svelte-h="svelte-4rqyu9">Mask generation is the task of generating semantically meaningful masks for an image. | |
| This task is very similar to <a href="semantic_segmentation">image segmentation</a>, but many differences exist. Image segmentation models are trained on labeled datasets and are limited to the classes they have seen during training; they return a set of masks and corresponding classes, given an image.</p> <p data-svelte-h="svelte-q9zsvz">Mask generation models are trained on large amounts of data and operate in two modes.</p> <ul data-svelte-h="svelte-1kcyijr"><li>Prompting mode: In this mode, the model takes in an image and a prompt, where a prompt can be a 2D point location (XY coordinates) in the image within an object or a bounding box surrounding an object. In prompting mode, the model only returns the mask over the object | |
| that the prompt is pointing out.</li> <li>Segment Everything mode: In segment everything, given an image, the model generates every mask in the image. To do so, a grid of points is generated and overlaid on the image for inference.</li> <li>Video Inference: The model accepts a video, and a point or box prompt in a video frame, which is tracked throughout the video. You can get more information on how to do video inference by following <a href="../model_doc/sam2">SAM 2 docs</a>.</li></ul> <p data-svelte-h="svelte-1fb4ueb">Mask generation task is supported by <a href="../model_doc/sam">Segment Anything Model (SAM)</a> and <a href="../model_doc/sam2">Segment Anything Model 2 (SAM2)</a>, while video inference is supported by <a href="../model_doc/sam2">Segment Anything Model 2 (SAM2)</a>. SAM is a powerful model that consists of a Vision Transformer-based image encoder, a prompt encoder, and a two-way transformer mask decoder. Images and prompts are encoded, and the decoder takes these embeddings and generates valid masks. Meanwhile, SAM 2 extends SAM by adding a memory module to track the masks.</p> <div class="flex justify-center" data-svelte-h="svelte-18bgtyc"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/sam.png" alt="SAM Architecture"></div> <p data-svelte-h="svelte-1ji6167">SAM serves as a powerful foundation model for segmentation as it has large data coverage. It is trained on | |
| <a href="https://ai.meta.com/datasets/segment-anything/" rel="nofollow">SA-1B</a>, a dataset with 1 million images and 1.1 billion masks.</p> <p data-svelte-h="svelte-1xy9go1">In this guide, you will learn how to:</p> <ul data-svelte-h="svelte-1xdcdcm"><li>Infer in segment everything mode with batching,</li> <li>Infer in point prompting mode,</li> <li>Infer in box prompting mode.</li></ul> <p data-svelte-h="svelte-1qvzz2d">First, let’s install <code>transformers</code>:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-bash "><!-- HTML_TAG_START -->pip install -q transformers<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="mask-generation-pipeline" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#mask-generation-pipeline"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Mask Generation Pipeline</span></h2> <p data-svelte-h="svelte-hy99c4">The easiest way to infer mask generation models is to use the <code>mask-generation</code> pipeline.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-meta">>>> </span><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> pipeline | |
| <span class="hljs-meta">>>> </span>checkpoint = <span class="hljs-string">"facebook/sam2-hiera-base-plus"</span> | |
| <span class="hljs-meta">>>> </span>mask_generator = pipeline(model=checkpoint, task=<span class="hljs-string">"mask-generation"</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-13rxk6q">Let’s see the image.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> PIL <span class="hljs-keyword">import</span> Image | |
| <span class="hljs-keyword">import</span> requests | |
| img_url = <span class="hljs-string">"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"</span> | |
| image = Image.<span class="hljs-built_in">open</span>(requests.get(img_url, stream=<span class="hljs-literal">True</span>).raw).convert(<span class="hljs-string">"RGB"</span>)<!-- HTML_TAG_END --></pre></div> <div class="flex justify-center" data-svelte-h="svelte-5ac7qo"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg" alt="Example Image"></div> <p data-svelte-h="svelte-xsc4t1">Let’s segment everything. <code>points-per-batch</code> enables parallel inference of points in segment everything mode. This enables faster inference, but consumes more memory. Moreover, SAM only enables batching over points and not the images. <code>pred_iou_thresh</code> is the IoU confidence threshold where only the masks above that certain threshold are returned.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->masks = mask_generator(image, points_per_batch=<span class="hljs-number">128</span>, pred_iou_thresh=<span class="hljs-number">0.88</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-rg5py">The <code>masks</code> looks like the following:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-bash "><!-- HTML_TAG_START -->{<span class="hljs-string">'masks'</span>: [tensor([[False, False, False, ..., True, True, True], | |
| [False, False, False, ..., True, True, True], | |
| [False, False, False, ..., True, True, True], | |
| ..., | |
| [False, False, False, ..., False, False, False], .. | |
| <span class="hljs-string">'scores'</span>: tensor([0.9874, 0.9793, 0.9780, 0.9776, ... 0.9016])}<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1kbs52p">We can visualize them like this:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt | |
| plt.imshow(image, cmap=<span class="hljs-string">'gray'</span>) | |
| <span class="hljs-keyword">for</span> i, mask <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(masks[<span class="hljs-string">"masks"</span>]): | |
| plt.imshow(mask, cmap=<span class="hljs-string">'viridis'</span>, alpha=<span class="hljs-number">0.1</span>, vmin=<span class="hljs-number">0</span>, vmax=<span class="hljs-number">1</span>) | |
| plt.axis(<span class="hljs-string">'off'</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-b8wd3m">Below is the original image in grayscale with colorful maps overlaid. Very impressive.</p> <div class="flex justify-center" data-svelte-h="svelte-1xv3qg"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee_segmented.png" alt="Visualized"></div> <h2 class="relative group"><a id="model-inference" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#model-inference"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Model Inference</span></h2> <h3 class="relative group"><a id="point-prompting" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#point-prompting"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Point Prompting</span></h3> <p data-svelte-h="svelte-kmstl8">You can also use the model without the pipeline. To do so, initialize the model and | |
| the processor.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> SamModel, SamProcessor | |
| <span class="hljs-keyword">from</span> accelerate <span class="hljs-keyword">import</span> Accelerator | |
| <span class="hljs-keyword">import</span> torch | |
| device = Accelerator().device | |
| model = SamModel.from_pretrained(<span class="hljs-string">"facebook/sam-vit-base"</span>).to(device) | |
| processor = SamProcessor.from_pretrained(<span class="hljs-string">"facebook/sam-vit-base"</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1mls8gr">To do point prompting, pass the input point to the processor, then take the processor output | |
| and pass it to the model for inference. To post-process the model output, pass the outputs and | |
| <code>original_sizes</code> are taken from the processor’s initial output. We need to pass these | |
| since the processor resizes the image, and the output needs to be extrapolated.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->input_points = [[[<span class="hljs-number">2592</span>, <span class="hljs-number">1728</span>]]] <span class="hljs-comment"># point location of the bee</span> | |
| inputs = processor(image, input_points=input_points, return_tensors=<span class="hljs-string">"pt"</span>).to(device) | |
| <span class="hljs-keyword">with</span> torch.no_grad(): | |
| outputs = model(**inputs) | |
| masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs[<span class="hljs-string">"original_sizes"</span>].cpu())<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-zoqnl5">We can visualize the three masks in the <code>masks</code> output.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt | |
| <span class="hljs-keyword">import</span> numpy <span class="hljs-keyword">as</span> np | |
| fig, axes = plt.subplots(<span class="hljs-number">1</span>, <span class="hljs-number">4</span>, figsize=(<span class="hljs-number">15</span>, <span class="hljs-number">5</span>)) | |
| axes[<span class="hljs-number">0</span>].imshow(image) | |
| axes[<span class="hljs-number">0</span>].set_title(<span class="hljs-string">'Original Image'</span>) | |
| mask_list = [masks[<span class="hljs-number">0</span>][<span class="hljs-number">0</span>][<span class="hljs-number">0</span>].numpy(), masks[<span class="hljs-number">0</span>][<span class="hljs-number">0</span>][<span class="hljs-number">1</span>].numpy(), masks[<span class="hljs-number">0</span>][<span class="hljs-number">0</span>][<span class="hljs-number">2</span>].numpy()] | |
| <span class="hljs-keyword">for</span> i, mask <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(mask_list, start=<span class="hljs-number">1</span>): | |
| overlayed_image = np.array(image).copy() | |
| overlayed_image[:,:,<span class="hljs-number">0</span>] = np.where(mask == <span class="hljs-number">1</span>, <span class="hljs-number">255</span>, overlayed_image[:,:,<span class="hljs-number">0</span>]) | |
| overlayed_image[:,:,<span class="hljs-number">1</span>] = np.where(mask == <span class="hljs-number">1</span>, <span class="hljs-number">0</span>, overlayed_image[:,:,<span class="hljs-number">1</span>]) | |
| overlayed_image[:,:,<span class="hljs-number">2</span>] = np.where(mask == <span class="hljs-number">1</span>, <span class="hljs-number">0</span>, overlayed_image[:,:,<span class="hljs-number">2</span>]) | |
| axes[i].imshow(overlayed_image) | |
| axes[i].set_title(<span class="hljs-string">f'Mask <span class="hljs-subst">{i}</span>'</span>) | |
| <span class="hljs-keyword">for</span> ax <span class="hljs-keyword">in</span> axes: | |
| ax.axis(<span class="hljs-string">'off'</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <div class="flex justify-center" data-svelte-h="svelte-od5m1y"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/masks.png" alt="Visualized"></div> <h3 class="relative group"><a id="box-prompting" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#box-prompting"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Box Prompting</span></h3> <p data-svelte-h="svelte-1hby767">You can also do box prompting in a similar fashion to point prompting. You can simply pass the input box in the format of a list | |
| <code>[x_min, y_min, x_max, y_max]</code> format along with the image to the <code>processor</code>. Take the processor output and directly pass it | |
| to the model, then post-process the output again.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-comment"># bounding box around the bee</span> | |
| box = [<span class="hljs-number">2350</span>, <span class="hljs-number">1600</span>, <span class="hljs-number">2850</span>, <span class="hljs-number">2100</span>] | |
| inputs = processor( | |
| image, | |
| input_boxes=[[[box]]], | |
| return_tensors=<span class="hljs-string">"pt"</span> | |
| ).to(model.device) | |
| <span class="hljs-keyword">with</span> torch.no_grad(): | |
| outputs = model(**inputs) | |
| mask = processor.image_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), | |
| inputs[<span class="hljs-string">"original_sizes"</span>].cpu(), | |
| )[<span class="hljs-number">0</span>][<span class="hljs-number">0</span>][<span class="hljs-number">0</span>].numpy()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-n8en77">You can visualize the bounding box around the bee as shown below.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.patches <span class="hljs-keyword">as</span> patches | |
| fig, ax = plt.subplots() | |
| ax.imshow(image) | |
| rectangle = patches.Rectangle((<span class="hljs-number">2350</span>, <span class="hljs-number">1600</span>), <span class="hljs-number">500</span>, <span class="hljs-number">500</span>, linewidth=<span class="hljs-number">2</span>, edgecolor=<span class="hljs-string">'r'</span>, facecolor=<span class="hljs-string">'none'</span>) | |
| ax.add_patch(rectangle) | |
| ax.axis(<span class="hljs-string">"off"</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <div class="flex justify-center" data-svelte-h="svelte-12t8f97"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/bbox.png" alt="Visualized Bbox"></div> <p data-svelte-h="svelte-vujakk">You can see the inference output below.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->fig, ax = plt.subplots() | |
| ax.imshow(image) | |
| ax.imshow(mask, cmap=<span class="hljs-string">'viridis'</span>, alpha=<span class="hljs-number">0.4</span>) | |
| ax.axis(<span class="hljs-string">"off"</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <div class="flex justify-center" data-svelte-h="svelte-1inger9"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/box_inference.png" alt="Visualized Inference"></div> <h2 class="relative group"><a id="fine-tuning-for-mask-generation" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#fine-tuning-for-mask-generation"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Fine-tuning for Mask Generation</span></h2> <p data-svelte-h="svelte-1pp2syj">We will fine-tune SAM2.1 on small part of MicroMat dataset for image matting. We need to install the <a href="https://github.com/Project-MONAI/MONAI" rel="nofollow">monai</a> library to use DICE loss, and <a href="https://huggingface.co/docs/trackio/index" rel="nofollow">trackio</a> for logging the masks during training.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-bash "><!-- HTML_TAG_START -->pip install -q datasets monai trackio | |
| ``` | |
| We can now load our dataset and take a look. | |
| ```python | |
| from datasets import load_dataset | |
| dataset = load_dataset(<span class="hljs-string">"merve/MicroMat-mini"</span>, <span class="hljs-built_in">split</span>=<span class="hljs-string">"train"</span>) | |
| dataset | |
| <span class="hljs-comment"># Dataset({</span> | |
| <span class="hljs-comment"># features: ['image', 'mask', 'prompt', 'image_id', 'object_id', 'sample_idx', 'granularity', </span> | |
| <span class="hljs-comment"># 'image_path', 'mask_path', 'prompt_path'], num_rows: 94</span> | |
| <span class="hljs-comment">#})</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-j228bb">We need image, mask and prompt columns. We split for train and test.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->dataset = dataset.train_test_split(test_size=<span class="hljs-number">0.1</span>) | |
| train_ds = dataset[<span class="hljs-string">"train"</span>] | |
| val_ds = dataset[<span class="hljs-string">"test"</span>]<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-9q1iss">Let’s take a look at a sample.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->train_ds[<span class="hljs-number">0</span>]<!-- HTML_TAG_END --></pre></div> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-json "><!-- HTML_TAG_START --> <span class="hljs-punctuation">{</span>'image'<span class="hljs-punctuation">:</span> <PIL.PngImagePlugin.PngImageFile image mode=RGB size=<span class="hljs-number">2040</span>x1356><span class="hljs-punctuation">,</span> | |
| 'mask'<span class="hljs-punctuation">:</span> <PIL.PngImagePlugin.PngImageFile image mode=L size=<span class="hljs-number">2040</span>x1356><span class="hljs-punctuation">,</span> | |
| 'prompt'<span class="hljs-punctuation">:</span> '<span class="hljs-punctuation">{</span><span class="hljs-attr">"point"</span><span class="hljs-punctuation">:</span> <span class="hljs-punctuation">[</span><span class="hljs-punctuation">[</span><span class="hljs-number">137</span><span class="hljs-punctuation">,</span> <span class="hljs-number">1165</span><span class="hljs-punctuation">,</span> <span class="hljs-number">1</span><span class="hljs-punctuation">]</span><span class="hljs-punctuation">,</span> <span class="hljs-punctuation">[</span><span class="hljs-number">77</span><span class="hljs-punctuation">,</span> <span class="hljs-number">1273</span><span class="hljs-punctuation">,</span> <span class="hljs-number">0</span><span class="hljs-punctuation">]</span><span class="hljs-punctuation">,</span> <span class="hljs-punctuation">[</span><span class="hljs-number">58</span><span class="hljs-punctuation">,</span> <span class="hljs-number">1351</span><span class="hljs-punctuation">,</span> <span class="hljs-number">0</span><span class="hljs-punctuation">]</span><span class="hljs-punctuation">]</span><span class="hljs-punctuation">,</span> <span class="hljs-attr">"bbox"</span><span class="hljs-punctuation">:</span> <span class="hljs-punctuation">[</span><span class="hljs-number">0</span><span class="hljs-punctuation">,</span> <span class="hljs-number">701</span><span class="hljs-punctuation">,</span> <span class="hljs-number">251</span><span class="hljs-punctuation">,</span> <span class="hljs-number">1356</span><span class="hljs-punctuation">]</span><span class="hljs-punctuation">}</span>'<span class="hljs-punctuation">,</span> | |
| 'image_id'<span class="hljs-punctuation">:</span> '<span class="hljs-number">0034</span>'<span class="hljs-punctuation">,</span> | |
| 'object_id'<span class="hljs-punctuation">:</span> '<span class="hljs-number">34</span>'<span class="hljs-punctuation">,</span> | |
| 'sample_idx'<span class="hljs-punctuation">:</span> <span class="hljs-number">1</span><span class="hljs-punctuation">,</span> | |
| 'granularity'<span class="hljs-punctuation">:</span> 'fine'<span class="hljs-punctuation">,</span> | |
| 'image_path'<span class="hljs-punctuation">:</span> '/content/MicroMat-mini/img/<span class="hljs-number">0034.</span>png'<span class="hljs-punctuation">,</span> | |
| 'mask_path'<span class="hljs-punctuation">:</span> '/content/MicroMat-mini/mask/<span class="hljs-number">0034</span>_34.png'<span class="hljs-punctuation">,</span> | |
| 'prompt_path'<span class="hljs-punctuation">:</span> '/content/MicroMat-mini/prompt/<span class="hljs-number">0034</span>_34.json'<span class="hljs-punctuation">}</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-17hozjg">Prompts are string of dictionaries, so you can get the bounding boxes as shown below.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> json | |
| json.loads(train_ds[<span class="hljs-string">"prompt"</span>][<span class="hljs-number">0</span>])[<span class="hljs-string">"bbox"</span>] | |
| <span class="hljs-comment"># [0, 701, 251, 1356]</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-15p5hot">Visualize an example image, prompt and mask.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt | |
| <span class="hljs-keyword">import</span> numpy <span class="hljs-keyword">as</span> np | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">show_mask</span>(<span class="hljs-params">mask, ax</span>): | |
| color = np.array([<span class="hljs-number">0.12</span>, <span class="hljs-number">0.56</span>, <span class="hljs-number">1.0</span>, <span class="hljs-number">0.6</span>]) | |
| mask = np.array(mask) | |
| h, w = mask.shape | |
| mask_image = mask.reshape(h, w, <span class="hljs-number">1</span>) * color.reshape(<span class="hljs-number">1</span>, <span class="hljs-number">1</span>, <span class="hljs-number">4</span>) | |
| ax.imshow(mask_image) | |
| x0, y0, x1, y1 = <span class="hljs-built_in">eval</span>(train_ds[<span class="hljs-string">"prompt"</span>][<span class="hljs-number">0</span>])[<span class="hljs-string">"bbox"</span>] | |
| ax.add_patch( | |
| plt.Rectangle((x0, y0), x1 - x0, y1 - y0, | |
| fill=<span class="hljs-literal">False</span>, edgecolor=<span class="hljs-string">"lime"</span>, linewidth=<span class="hljs-number">2</span>)) | |
| example = train_ds[<span class="hljs-number">0</span>] | |
| image = np.array(example[<span class="hljs-string">"image"</span>]) | |
| ground_truth_mask = np.array(example[<span class="hljs-string">"mask"</span>]) | |
| fig, ax = plt.subplots() | |
| ax.imshow(image) | |
| show_mask(ground_truth_mask, ax) | |
| ax.set_title(<span class="hljs-string">"Ground truth mask"</span>) | |
| ax.set_axis_off() | |
| plt.show() <!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1yopd03">Now we can define our dataset for loading the data. SAMDataset wraps our dataset and formats each sample the way the SAM processor expects. So instead of raw images and masks, you get processed images, bounding boxes, and ground-truth masks ready for training.</p> <p data-svelte-h="svelte-1iv3hz2">By default, processor resizes images, so on top of images and masks, it also returns original sizes. We also need to binarize the mask as it has values [0, 255].</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> torch.utils.data <span class="hljs-keyword">import</span> Dataset | |
| <span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">class</span> <span class="hljs-title class_">SAMDataset</span>(<span class="hljs-title class_ inherited__">Dataset</span>): | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, dataset, processor</span>): | |
| self.dataset = dataset | |
| self.processor = processor | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">__len__</span>(<span class="hljs-params">self</span>): | |
| <span class="hljs-keyword">return</span> <span class="hljs-built_in">len</span>(self.dataset) | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">__getitem__</span>(<span class="hljs-params">self, idx</span>): | |
| item = self.dataset[idx] | |
| image = item[<span class="hljs-string">"image"</span>] | |
| prompt = <span class="hljs-built_in">eval</span>(item[<span class="hljs-string">"prompt"</span>])[<span class="hljs-string">"bbox"</span>] | |
| inputs = self.processor(image, input_boxes=[[prompt]], return_tensors=<span class="hljs-string">"pt"</span>) | |
| inputs[<span class="hljs-string">"ground_truth_mask"</span>] = (np.array(item[<span class="hljs-string">"mask"</span>]) > <span class="hljs-number">0</span>).astype(np.float32) | |
| inputs[<span class="hljs-string">"original_image_size"</span>] = torch.tensor(image.size[::-<span class="hljs-number">1</span>]) | |
| <span class="hljs-keyword">return</span> inputs<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-elhaxu">We can initialize the processor and the dataset with it.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> Sam2Processor | |
| processor = Sam2Processor.from_pretrained(<span class="hljs-string">"facebook/sam2.1-hiera-small"</span>) | |
| train_dataset = SAMDataset(dataset=train_ds, processor=processor) | |
| ``` | |
| We need to define a data collator that will turn varying size of ground truth masks to batches of reshaped masks <span class="hljs-keyword">in</span> same shape. We reshape them using nearest neighbor interpolation. We also make batched tensors <span class="hljs-keyword">for</span> rest of the elements <span class="hljs-keyword">in</span> the batch. If your masks are <span class="hljs-built_in">all</span> of same size, feel free to skip this step. | |
| ```python | |
| <span class="hljs-keyword">import</span> torch.nn.functional <span class="hljs-keyword">as</span> F | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">collate_fn</span>(<span class="hljs-params">batch, target_hw=(<span class="hljs-params"><span class="hljs-number">256</span>, <span class="hljs-number">256</span></span>)</span>): | |
| pixel_values = torch.cat([item[<span class="hljs-string">"pixel_values"</span>] <span class="hljs-keyword">for</span> item <span class="hljs-keyword">in</span> batch], dim=<span class="hljs-number">0</span>) | |
| original_sizes = torch.stack([item[<span class="hljs-string">"original_sizes"</span>] <span class="hljs-keyword">for</span> item <span class="hljs-keyword">in</span> batch]) | |
| input_boxes = torch.cat([item[<span class="hljs-string">"input_boxes"</span>] <span class="hljs-keyword">for</span> item <span class="hljs-keyword">in</span> batch], dim=<span class="hljs-number">0</span>) | |
| ground_truth_masks = torch.cat([ | |
| F.interpolate( | |
| torch.as_tensor(x[<span class="hljs-string">"ground_truth_mask"</span>]).unsqueeze(<span class="hljs-number">0</span>).unsqueeze(<span class="hljs-number">0</span>).<span class="hljs-built_in">float</span>(), | |
| size=(<span class="hljs-number">256</span>, <span class="hljs-number">256</span>), | |
| mode=<span class="hljs-string">"nearest"</span> | |
| ) | |
| <span class="hljs-keyword">for</span> x <span class="hljs-keyword">in</span> batch | |
| ], dim=<span class="hljs-number">0</span>).long() | |
| <span class="hljs-keyword">return</span> { | |
| <span class="hljs-string">"pixel_values"</span>: pixel_values, | |
| <span class="hljs-string">"original_sizes"</span>: original_sizes, | |
| <span class="hljs-string">"input_boxes"</span>: input_boxes, | |
| <span class="hljs-string">"ground_truth_mask"</span>: ground_truth_masks, | |
| <span class="hljs-string">"original_image_size"</span>: torch.stack([item[<span class="hljs-string">"original_image_size"</span>] <span class="hljs-keyword">for</span> item <span class="hljs-keyword">in</span> batch]), | |
| } | |
| <span class="hljs-keyword">from</span> torch.utils.data <span class="hljs-keyword">import</span> DataLoader | |
| train_dataloader = DataLoader( | |
| train_dataset, | |
| batch_size=<span class="hljs-number">4</span>, | |
| shuffle=<span class="hljs-literal">True</span>, | |
| collate_fn=collate_fn, | |
| )<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1jchvz">Let’s take a look at what the data loader yields.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->batch = <span class="hljs-built_in">next</span>(<span class="hljs-built_in">iter</span>(train_dataloader)) | |
| <span class="hljs-keyword">for</span> k,v <span class="hljs-keyword">in</span> batch.items(): | |
| <span class="hljs-built_in">print</span>(k,v.shape) | |
| <span class="hljs-comment"># pixel_values torch.Size([4, 3, 1024, 1024])</span> | |
| <span class="hljs-comment"># original_sizes torch.Size([4, 1, 2])</span> | |
| <span class="hljs-comment"># input_boxes torch.Size([4, 1, 4])</span> | |
| <span class="hljs-comment"># ground_truth_mask torch.Size([4, 1, 256, 256])</span> | |
| <span class="hljs-comment">#original_image_size torch.Size([4, 2])</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1l5sx5f">We will now load the model and freeze the vision and the prompt encoder to only train the mask decoder.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> Sam2Model | |
| model = Sam2Model.from_pretrained(<span class="hljs-string">"facebook/sam2.1-hiera-small"</span>) | |
| <span class="hljs-keyword">for</span> name, param <span class="hljs-keyword">in</span> model.named_parameters(): | |
| <span class="hljs-keyword">if</span> name.startswith(<span class="hljs-string">"vision_encoder"</span>) <span class="hljs-keyword">or</span> name.startswith(<span class="hljs-string">"prompt_encoder"</span>): | |
| param.requires_grad_(<span class="hljs-literal">False</span>) | |
| ``` | |
| We can now define the optimizer <span class="hljs-keyword">and</span> the loss function. | |
| ```python | |
| <span class="hljs-keyword">from</span> torch.optim <span class="hljs-keyword">import</span> Adam | |
| <span class="hljs-keyword">import</span> monai | |
| optimizer = Adam(model.mask_decoder.parameters(), lr=<span class="hljs-number">1e-5</span>, weight_decay=<span class="hljs-number">0</span>) | |
| seg_loss = monai.losses.DiceCELoss(sigmoid=<span class="hljs-literal">True</span>, squared_pred=<span class="hljs-literal">True</span>, reduction=<span class="hljs-string">'mean'</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1jpsdv7">Let’s see how the model performs before training.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt | |
| item = val_ds[<span class="hljs-number">1</span>] | |
| img = item[<span class="hljs-string">"image"</span>] | |
| bbox = json.loads(item[<span class="hljs-string">"prompt"</span>])[<span class="hljs-string">"bbox"</span>] | |
| inputs = processor(images=img, input_boxes=[[bbox]], return_tensors=<span class="hljs-string">"pt"</span>).to(model.device) | |
| <span class="hljs-keyword">with</span> torch.no_grad(): | |
| outputs = model(**inputs) | |
| masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs[<span class="hljs-string">"original_sizes"</span>])[<span class="hljs-number">0</span>] | |
| preds = masks.squeeze(<span class="hljs-number">0</span>) | |
| mask = (preds[<span class="hljs-number">0</span>] > <span class="hljs-number">0</span>).cpu().numpy() | |
| overlay = np.asarray(img, dtype=np.uint8).copy() | |
| overlay[mask] = <span class="hljs-number">0.55</span> * overlay[mask] + <span class="hljs-number">0.45</span> * np.array([<span class="hljs-number">0</span>, <span class="hljs-number">255</span>, <span class="hljs-number">0</span>], dtype=np.float32) | |
| plt.imshow(overlay) | |
| plt.axis(<span class="hljs-string">"off"</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-xruqcd"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/sam2_before_training.png" alt="SAM2 result after training"></p> <p data-svelte-h="svelte-ks36k4">We need to log our predictions to trackio so we can monitor the model improvement in the middle of the training.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> PIL <span class="hljs-keyword">import</span> Image | |
| <span class="hljs-keyword">import</span> trackio | |
| <span class="hljs-keyword">import</span> json | |
| <span class="hljs-meta">@torch.no_grad()</span> | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">predict_fn</span>(<span class="hljs-params">img, bbox</span>): | |
| inputs = processor(images=img, input_boxes=[[bbox]], return_tensors=<span class="hljs-string">"pt"</span>).to(model.device) | |
| <span class="hljs-keyword">with</span> torch.no_grad(): | |
| outputs = model(**inputs) | |
| masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs[<span class="hljs-string">"original_sizes"</span>])[<span class="hljs-number">0</span>] | |
| <span class="hljs-keyword">return</span> masks | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">log_eval_masks_trackio</span>(<span class="hljs-params">dataset, indices, step, predict_fn, project=<span class="hljs-literal">None</span>, sample_cap=<span class="hljs-number">8</span></span>): | |
| logs = {<span class="hljs-string">"eval/step"</span>: <span class="hljs-built_in">int</span>(step)} | |
| <span class="hljs-keyword">for</span> idx <span class="hljs-keyword">in</span> indices[:sample_cap]: | |
| item = dataset[idx] | |
| img = item[<span class="hljs-string">"image"</span>] | |
| bbox = json.loads(item[<span class="hljs-string">"prompt"</span>])[<span class="hljs-string">"bbox"</span>] | |
| preds = predict_fn(img, bbox) | |
| preds = preds.squeeze(<span class="hljs-number">0</span>) | |
| mask = (preds[<span class="hljs-number">0</span>] > <span class="hljs-number">0</span>).cpu().numpy() | |
| overlay = np.asarray(img, dtype=np.uint8).copy() | |
| overlay[mask] = <span class="hljs-number">0.55</span> * overlay[mask] + <span class="hljs-number">0.45</span> * np.array([<span class="hljs-number">0</span>, <span class="hljs-number">255</span>, <span class="hljs-number">0</span>], dtype=np.float32) | |
| logs[<span class="hljs-string">f"<span class="hljs-subst">{idx}</span>/overlay"</span>] = trackio.Image(overlay, caption=<span class="hljs-string">"overlay"</span>) | |
| trackio.log(logs)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-e45dkf">We can now write our training loop and train!</p> <p data-svelte-h="svelte-1tltiio">Notice how we log our loss and evaluation masks with trackio.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> tqdm <span class="hljs-keyword">import</span> tqdm | |
| <span class="hljs-keyword">from</span> statistics <span class="hljs-keyword">import</span> mean | |
| <span class="hljs-keyword">import</span> trackio | |
| <span class="hljs-keyword">import</span> torch | |
| num_epochs = <span class="hljs-number">30</span> | |
| device = <span class="hljs-string">"cuda"</span> <span class="hljs-keyword">if</span> torch.cuda.is_available() <span class="hljs-keyword">else</span> <span class="hljs-string">"cpu"</span> | |
| model.to(device) | |
| model.train() | |
| trackio.init(project=<span class="hljs-string">"mask-eval"</span>) | |
| <span class="hljs-keyword">for</span> epoch <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(num_epochs): | |
| epoch_losses = [] | |
| <span class="hljs-keyword">for</span> batch <span class="hljs-keyword">in</span> tqdm(train_dataloader): | |
| outputs = model(pixel_values=batch[<span class="hljs-string">"pixel_values"</span>].to(device), | |
| input_boxes=batch[<span class="hljs-string">"input_boxes"</span>].to(device), | |
| multimask_output=<span class="hljs-literal">False</span>) | |
| predicted_masks = outputs.pred_masks.squeeze(<span class="hljs-number">1</span>) | |
| ground_truth_masks = batch[<span class="hljs-string">"ground_truth_mask"</span>].<span class="hljs-built_in">float</span>().to(device) | |
| loss = seg_loss(predicted_masks, ground_truth_masks) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| epoch_losses.append(loss.item()) | |
| log_eval_masks_trackio(dataset=val_ds, indices=[<span class="hljs-number">0</span>, <span class="hljs-number">3</span>, <span class="hljs-number">6</span>, <span class="hljs-number">9</span>], step=epoch, predict_fn=predict_fn, project=<span class="hljs-string">"mask-eval"</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f'Epoch: <span class="hljs-subst">{epoch}</span>'</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f'Mean loss: <span class="hljs-subst">{mean(epoch_losses)}</span>'</span>) | |
| trackio.log({<span class="hljs-string">"loss"</span>: mean(epoch_losses)}) | |
| trackio.finish()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1sznapg">Let’s put the trained model to test.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> matplotlib.pyplot <span class="hljs-keyword">as</span> plt | |
| item = val_ds[<span class="hljs-number">1</span>] | |
| img = item[<span class="hljs-string">"image"</span>] | |
| bbox = json.loads(item[<span class="hljs-string">"prompt"</span>])[<span class="hljs-string">"bbox"</span>] | |
| inputs = processor(images=img, input_boxes=[[bbox]], return_tensors=<span class="hljs-string">"pt"</span>).to(model.device) | |
| <span class="hljs-keyword">with</span> torch.no_grad(): | |
| outputs = model(**inputs) | |
| preds = processor.post_process_masks(outputs.pred_masks.cpu(), inputs[<span class="hljs-string">"original_sizes"</span>])[<span class="hljs-number">0</span>] | |
| preds = preds.squeeze(<span class="hljs-number">0</span>) | |
| mask = (preds[<span class="hljs-number">0</span>] > <span class="hljs-number">0</span>).cpu().numpy() | |
| overlay = np.asarray(img, dtype=np.uint8).copy() | |
| overlay[mask] = <span class="hljs-number">0.55</span> * overlay[mask] + <span class="hljs-number">0.45</span> * np.array([<span class="hljs-number">0</span>, <span class="hljs-number">255</span>, <span class="hljs-number">0</span>], dtype=np.float32) | |
| plt.imshow(overlay) | |
| plt.axis(<span class="hljs-string">"off"</span>) | |
| plt.show()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1lujldg">Great improvement after only training for 20 epochs on a small dataset!</p> <p data-svelte-h="svelte-1xyqtii"><img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/sam2_after_training.png" alt="SAM2 result after training"></p> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/transformers/blob/main/docs/source/en/tasks/mask_generation.md" target="_blank"><svg class="mr-1" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M31,16l-7,7l-1.41-1.41L28.17,16l-5.58-5.59L24,9l7,7z"></path><path d="M1,16l7-7l1.41,1.41L3.83,16l5.58,5.59L8,23l-7-7z"></path><path d="M12.419,25.484L17.639,6.552l1.932,0.518L14.351,26.002z"></path></svg> <span data-svelte-h="svelte-zjs2n5"><span class="underline">Update</span> on GitHub</span></a> <p></p> | |
| <script> | |
| { | |
| __sveltekit_sw4oup = { | |
| assets: "/docs/transformers/pr_39895/en", | |
| base: "/docs/transformers/pr_39895/en", | |
| env: {} | |
| }; | |
| const element = document.currentScript.parentElement; | |
| const data = [null,null]; | |
| Promise.all([ | |
| import("/docs/transformers/pr_39895/en/_app/immutable/entry/start.ed44d8d9.js"), | |
| import("/docs/transformers/pr_39895/en/_app/immutable/entry/app.75a9946c.js") | |
| ]).then(([kit, app]) => { | |
| kit.start(app, element, { | |
| node_ids: [0, 680], | |
| data, | |
| form: null, | |
| error: null | |
| }); | |
| }); | |
| } | |
| </script> | |
Xet Storage Details
- Size:
- 81.5 kB
- Xet hash:
- 7ed8bd99547ece724c9299bbf6da3d0ce8d439f8d837beb224cfb0356bf0cfd7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.