CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/vision/sam/" />  <meta property="og:title" content="Keras documentation: Segment Anything Model with 🤗Transformers"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Segment Anything Model with 🤗Transformers"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Segment Anything Model with 🤗Transformers</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink active" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink2" href="/examples/vision/image_classification_from_scratch/">Image classification from scratch</a> <a class="nav-sublink2" href="/examples/vision/mnist_convnet/">Simple MNIST convnet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_efficientnet_fine_tuning/">Image classification via fine-tuning with EfficientNet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_with_vision_transformer/">Image classification with Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/attention_mil_classification/">Classification using Attention-based Deep Multiple Instance Learning</a> <a class="nav-sublink2" href="/examples/vision/mlp_image_classification/">Image classification with modern MLP models</a> <a class="nav-sublink2" href="/examples/vision/mobilevit/">A mobile-friendly Transformer-based model for image classification</a> <a class="nav-sublink2" href="/examples/vision/xray_classification_with_tpus/">Pneumonia Classification on TPU</a> <a class="nav-sublink2" href="/examples/vision/cct/">Compact Convolutional Transformers</a> <a class="nav-sublink2" href="/examples/vision/convmixer/">Image classification with ConvMixer</a> <a class="nav-sublink2" href="/examples/vision/eanet/">Image classification with EANet (External Attention Transformer)</a> <a class="nav-sublink2" href="/examples/vision/involution/">Involutional neural networks</a> <a class="nav-sublink2" href="/examples/vision/perceiver_image_classification/">Image classification with Perceiver</a> <a class="nav-sublink2" href="/examples/vision/reptile/">Few-Shot learning with Reptile</a> <a class="nav-sublink2" href="/examples/vision/semisupervised_simclr/">Semi-supervised image classification using contrastive pretraining with SimCLR</a> <a class="nav-sublink2" href="/examples/vision/swin_transformers/">Image classification with Swin Transformers</a> <a class="nav-sublink2" href="/examples/vision/vit_small_ds/">Train a Vision Transformer on small datasets</a> <a class="nav-sublink2" href="/examples/vision/shiftvit/">A Vision Transformer without Attention</a> <a class="nav-sublink2" href="/examples/vision/image_classification_using_global_context_vision_transformer/">Image Classification using Global Context Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/oxford_pets_image_segmentation/">Image segmentation with a U-Net-like architecture</a> <a class="nav-sublink2" href="/examples/vision/deeplabv3_plus/">Multiclass semantic segmentation using DeepLabV3+</a> <a class="nav-sublink2" href="/examples/vision/basnet_segmentation/">Highly accurate boundaries segmentation using BASNet</a> <a class="nav-sublink2" href="/examples/vision/fully_convolutional_network/">Image Segmentation using Composable Fully-Convolutional Networks</a> <a class="nav-sublink2" href="/examples/vision/retinanet/">Object Detection with RetinaNet</a> <a class="nav-sublink2" href="/examples/vision/keypoint_detection/">Keypoint Detection with Transfer Learning</a> <a class="nav-sublink2" href="/examples/vision/object_detection_using_vision_transformer/">Object detection with Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/3D_image_classification/">3D image classification from CT scans</a> <a class="nav-sublink2" href="/examples/vision/depth_estimation/">Monocular depth estimation</a> <a class="nav-sublink2" href="/examples/vision/nerf/">3D volumetric rendering with NeRF</a> <a class="nav-sublink2" href="/examples/vision/pointnet_segmentation/">Point cloud segmentation with PointNet</a> <a class="nav-sublink2" href="/examples/vision/pointnet/">Point cloud classification</a> <a class="nav-sublink2" href="/examples/vision/captcha_ocr/">OCR model for reading Captchas</a> <a class="nav-sublink2" href="/examples/vision/handwriting_recognition/">Handwriting recognition</a> <a class="nav-sublink2" href="/examples/vision/autoencoder/">Convolutional autoencoder for image denoising</a> <a class="nav-sublink2" href="/examples/vision/mirnet/">Low-light image enhancement using MIRNet</a> <a class="nav-sublink2" href="/examples/vision/super_resolution_sub_pixel/">Image Super-Resolution using an Efficient Sub-Pixel CNN</a> <a class="nav-sublink2" href="/examples/vision/edsr/">Enhanced Deep Residual Networks for single-image super-resolution</a> <a class="nav-sublink2" href="/examples/vision/zero_dce/">Zero-DCE for low-light image enhancement</a> <a class="nav-sublink2" href="/examples/vision/cutmix/">CutMix data augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/mixup/">MixUp augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/randaugment/">RandAugment for Image Classification for Improved Robustness</a> <a class="nav-sublink2" href="/examples/vision/image_captioning/">Image captioning</a> <a class="nav-sublink2" href="/examples/vision/nl_image_search/">Natural language image search with a Dual Encoder</a> <a class="nav-sublink2" href="/examples/vision/visualizing_what_convnets_learn/">Visualizing what convnets learn</a> <a class="nav-sublink2" href="/examples/vision/integrated_gradients/">Model interpretability with Integrated Gradients</a> <a class="nav-sublink2" href="/examples/vision/probing_vits/">Investigating Vision Transformer representations</a> <a class="nav-sublink2" href="/examples/vision/grad_cam/">Grad-CAM class activation visualization</a> <a class="nav-sublink2" href="/examples/vision/near_dup_search/">Near-duplicate image search</a> <a class="nav-sublink2" href="/examples/vision/semantic_image_clustering/">Semantic Image Clustering</a> <a class="nav-sublink2" href="/examples/vision/siamese_contrastive/">Image similarity estimation using a Siamese Network with a contrastive loss</a> <a class="nav-sublink2" href="/examples/vision/siamese_network/">Image similarity estimation using a Siamese Network with a triplet loss</a> <a class="nav-sublink2" href="/examples/vision/metric_learning/">Metric learning for image similarity search</a> <a class="nav-sublink2" href="/examples/vision/metric_learning_tf_similarity/">Metric learning for image similarity search using TensorFlow Similarity</a> <a class="nav-sublink2" href="/examples/vision/nnclr/">Self-supervised contrastive learning with NNCLR</a> <a class="nav-sublink2" href="/examples/vision/video_classification/">Video Classification with a CNN-RNN Architecture</a> <a class="nav-sublink2" href="/examples/vision/conv_lstm/">Next-Frame Video Prediction with Convolutional LSTMs</a> <a class="nav-sublink2" href="/examples/vision/video_transformers/">Video Classification with Transformers</a> <a class="nav-sublink2" href="/examples/vision/vivit/">Video Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/bit/">Image Classification using BigTransfer (BiT)</a> <a class="nav-sublink2" href="/examples/vision/gradient_centralization/">Gradient Centralization for Better Training Performance</a> <a class="nav-sublink2" href="/examples/vision/token_learner/">Learning to tokenize in Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/knowledge_distillation/">Knowledge Distillation</a> <a class="nav-sublink2" href="/examples/vision/fixres/">FixRes: Fixing train-test resolution discrepancy</a> <a class="nav-sublink2" href="/examples/vision/cait/">Class Attention Image Transformers with LayerScale</a> <a class="nav-sublink2" href="/examples/vision/patch_convnet/">Augmenting convnets with aggregated attention</a> <a class="nav-sublink2" href="/examples/vision/learnable_resizer/">Learning to Resize</a> <a class="nav-sublink2" href="/examples/vision/adamatch/">Semi-supervision and domain adaptation with AdaMatch</a> <a class="nav-sublink2" href="/examples/vision/barlow_twins/">Barlow Twins for Contrastive SSL</a> <a class="nav-sublink2" href="/examples/vision/consistency_training/">Consistency training with supervision</a> <a class="nav-sublink2" href="/examples/vision/deit/">Distilling Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/focal_modulation_network/">Focal Modulation: A replacement for Self-Attention</a> <a class="nav-sublink2" href="/examples/vision/forwardforward/">Using the Forward-Forward Algorithm for Image Classification</a> <a class="nav-sublink2" href="/examples/vision/masked_image_modeling/">Masked image modeling with Autoencoders</a> <a class="nav-sublink2 active" href="/examples/vision/sam/">Segment Anything Model with 🤗Transformers</a> <a class="nav-sublink2" href="/examples/vision/segformer/">Semantic segmentation with SegFormer and Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/vision/simsiam/">Self-supervised contrastive learning with SimSiam</a> <a class="nav-sublink2" href="/examples/vision/supervised-contrastive-learning/">Supervised Contrastive Learning</a> <a class="nav-sublink2" href="/examples/vision/temporal_latent_bottleneck/">When Recurrence meets Transformers</a> <a class="nav-sublink2" href="/examples/vision/yolov8/">Efficient Object Detection with YOLOV8 and KerasCV</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparameter Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> <a class="nav-link" href="/keras_cv/" role="tab" aria-selected="">KerasCV: Computer Vision Workflows</a> <a class="nav-link" href="/keras_nlp/" role="tab" aria-selected="">KerasNLP: Natural Language Workflows</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/vision/'>Computer Vision</a> / Segment Anything Model with 🤗Transformers </div> <div class='k-content'> <h1 id="segment-anything-model-with-🤗transformers">Segment Anything Model with 🤗Transformers</h1> Authors: <a href="https://twitter.com/mervenoyann">Merve Noyan</a> & <a href="https://twitter.com/RisingSayak">Sayak Paul</a> Date created: 2023/07/11 Last modified: 2023/07/11 Description: Fine-tuning Segment Anything Model using Keras and 🤗 Transformers. <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/vision/ipynb/sam.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/vision/sam.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> Large language models (LLMs) make it easy for the end users to apply them to various applications through "prompting". For example if we wanted an LLM to predict the sentiment of the following sentence – "That movie was amazing, I thoroughly enjoyed it" – we'd do prompt the LLM with something like: <blockquote> What's the sentiment of the following sentence: "That movie was amazing, I thoroughly enjoyed it"? </blockquote> In return, the LLM would return sentiment token. But when it comes to visual recognition tasks, how can we engineer "visual" cues to prompt foundation vision models? For example, we could have an input image and prompt the model with bounding box on that image and ask it to perform segmentation. The bounding box would serve as our visual prompt here. In the <a href="https://segment-anything.com/">Segment Anything Model</a> (dubbed as SAM), researchers from Meta extended the space of language prompting to visual prompting. SAM is capable of performing zero-shot segmentation with a prompt input, inspired by large language models. The prompt here can be a set of foreground/background points, free text, a box or a mask. There are many downstream segmentation tasks, including semantic segmentation and edge detection. The goal of SAM is to enable all of these downstream segmentation tasks through prompting. In this example, we'll learn how to use the SAM model from 🤗 Transformers for performing inference and fine-tuning. <hr /> <h2 id="installation">Installation</h2> <div class="codehilite"><pre><code>!!pip install -q git+https://github.com/huggingface/transformers </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[] </code></pre></div> </div> Let's import everything we need for this example. <div class="codehilite"><pre><code>from tensorflow import keras from transformers import TFSamModel, SamProcessor import numpy as np import tensorflow as tf import matplotlib.pyplot as plt from tensorflow.python.ops.numpy_ops import np_config from PIL import Image import requests import glob import os </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>/Users/mervenoyan/miniforge3/envs/py310/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm </code></pre></div> </div> <hr /> <h2 id="sam-in-a-few-words">SAM in a few words</h2> SAM has the following components: <table> <thead> <tr> <th style="text-align: center;"><img alt="" src="https://imgur.com/oLfdwuB.png" /></th> </tr> </thead> <tbody> <tr> <td style="text-align: center;">Image taken from the official <a href="https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/">SAM blog post</a></td> </tr> </tbody> </table> The image encoder is responsible for computing image embeddings. When interacting with SAM, we compute the image embedding one time (as the image encoder is heavy) and then reuse it with different prompts mentioned above (points, bounding boxes, masks). Points and boxes (so-called sparse prompts) go through a lightweight prompt encoder, while masks (dense prompts) go through a convolutional layer. We couple the image embedding extracted from the image encoder and the prompt embedding and both go to a lightweight mask decoder. The decoder is responsible for predicting the mask. <table> <thead> <tr> <th style="text-align: center;"><img alt="" src="https://i.imgur.com/QQ9Ts5T.png" /></th> </tr> </thead> <tbody> <tr> <td style="text-align: center;">Figure taken from the <a href="https://arxiv.org/abs/2304.02643">SAM paper</a></td> </tr> </tbody> </table> SAM was pre-trained to predict a valid mask for any acceptable prompt. This requirement allows SAM to output a valid mask even when the prompt is ambiguous to understand – this makes SAM ambiguity-aware. Moreover, SAM predicts multiple masks for a single prompt. We highly encourage you to check out the <a href="https://arxiv.org/abs/2304.02643">SAM paper</a> and the <a href="https://ai.facebook.com/blog/segment-anything-foundation-model-image-segmentation/">blog post</a> to learn more about the additional details of SAM and the dataset used to pre-trained it. <hr /> <h2 id="running-inference-with-sam">Running inference with SAM</h2> There are three checkpoints for SAM: <ul> <li><a href="https://huggingface.co/facebook/sam-vit-base">sam-vit-base</a></li> <li><a href="https://huggingface.co/facebook/sam-vit-large">sam-vit-large</a></li> <li><a href="https://huggingface.co/facebook/sam-vit-huge">sam-vit-huge</a>.</li> </ul> We load <code>sam-vit-base</code> in <a href="https://huggingface.co/docs/transformers/main/model_doc/sam#transformers.TFSamModel"><code>TFSamModel</code></a>. We also need <code>SamProcessor</code>for the associated checkpoint. <div class="codehilite"><pre><code>model = TFSamModel.from_pretrained("facebook/sam-vit-base") processor = SamProcessor.from_pretrained("facebook/sam-vit-base") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>All model checkpoint layers were used when initializing TFSamModel. </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>All the layers of TFSamModel were initialized from the model checkpoint at facebook/sam-vit-base. If your task is similar to the task the model of the checkpoint was trained on, you can already use TFSamModel for predictions without further training. </code></pre></div> </div> Next, we write some utility functions for visualization. Most of these functions are taken from <a href="https://github.com/huggingface/notebooks/blob/main/examples/segment_anything.ipynb">this notebook</a>. <div class="codehilite"><pre><code>np_config.enable_numpy_behavior() def show_mask(mask, ax, random_color=False): if random_color: color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) else: color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6]) h, w = mask.shape[-2:] mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) ax.imshow(mask_image) def show_box(box, ax): x0, y0 = box[0], box[1] w, h = box[2] - box[0], box[3] - box[1] ax.add_patch( plt.Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2) ) def show_boxes_on_image(raw_image, boxes): plt.figure(figsize=(10, 10)) plt.imshow(raw_image) for box in boxes: show_box(box, plt.gca()) plt.axis("on") plt.show() def show_points_on_image(raw_image, input_points, input_labels=None): plt.figure(figsize=(10, 10)) plt.imshow(raw_image) input_points = np.array(input_points) if input_labels is None: labels = np.ones_like(input_points[:, 0]) else: labels = np.array(input_labels) show_points(input_points, labels, plt.gca()) plt.axis("on") plt.show() def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None): plt.figure(figsize=(10, 10)) plt.imshow(raw_image) input_points = np.array(input_points) if input_labels is None: labels = np.ones_like(input_points[:, 0]) else: labels = np.array(input_labels) show_points(input_points, labels, plt.gca()) for box in boxes: show_box(box, plt.gca()) plt.axis("on") plt.show() def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None): plt.figure(figsize=(10, 10)) plt.imshow(raw_image) input_points = np.array(input_points) if input_labels is None: labels = np.ones_like(input_points[:, 0]) else: labels = np.array(input_labels) show_points(input_points, labels, plt.gca()) for box in boxes: show_box(box, plt.gca()) plt.axis("on") plt.show() def show_points(coords, labels, ax, marker_size=375): pos_points = coords[labels == 1] neg_points = coords[labels == 0] ax.scatter( pos_points[:, 0], pos_points[:, 1], color="green", marker="*", s=marker_size, edgecolor="white", linewidth=1.25, ) ax.scatter( neg_points[:, 0], neg_points[:, 1], color="red", marker="*", s=marker_size, edgecolor="white", linewidth=1.25, ) def show_masks_on_image(raw_image, masks, scores): if len(masks[0].shape) == 4: final_masks = tf.squeeze(masks[0]) if scores.shape[0] == 1: final_scores = tf.squeeze(scores) nb_predictions = scores.shape[-1] fig, axes = plt.subplots(1, nb_predictions, figsize=(15, 15)) for i, (mask, score) in enumerate(zip(final_masks, final_scores)): mask = tf.stop_gradient(mask) axes[i].imshow(np.array(raw_image)) show_mask(mask, axes[i]) axes[i].title.set_text(f"Mask {i+1}, Score: {score.numpy().item():.3f}") axes[i].axis("off") plt.show() </code></pre></div> We will segment a car image using a point prompt. Make sure to set <code>return_tensors</code> to <code>tf</code> when calling the processor. Let's load an image of a car and segment it. <div class="codehilite"><pre><code>img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png" raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") plt.imshow(raw_image) plt.show() </code></pre></div> <img alt="png" src="/img/examples/vision/sam/sam_14_0.png" /> Let's now define a set of points we will use as the prompt. <div class="codehilite"><pre><code>input_points = [[[450, 600]]] # Visualize a single point. show_points_on_image(raw_image, input_points[0]) </code></pre></div> <img alt="png" src="/img/examples/vision/sam/sam_16_0.png" /> And segment: <div class="codehilite"><pre><code># Preprocess the input image. inputs = processor(raw_image, input_points=input_points, return_tensors="tf") # Predict for segmentation with the prompt. outputs = model(**inputs) </code></pre></div> <code>outputs</code> has got two attributes of our interest: <ul> <li><code>outputs.pred_masks</code>: which denotes the predicted masks.</li> <li><code>outputs.iou_scores</code>: which denotes the IoU scores associated with the masks.</li> </ul> Let's post-process the masks and visualize them with their IoU scores: <div class="codehilite"><pre><code>masks = processor.image_processor.post_process_masks( outputs.pred_masks, inputs["original_sizes"], inputs["reshaped_input_sizes"], return_tensors="tf", ) show_masks_on_image(raw_image, masks, outputs.iou_scores) </code></pre></div> <img alt="png" src="/img/examples/vision/sam/sam_21_0.png" /> And there we go! As can be noticed, all the masks are valid masks for the point prompt we provided. SAM is flexible enough to support different visual prompts and we encourage you to check out <a href="https://github.com/huggingface/notebooks/blob/main/examples/segment_anything.ipynb">this notebook</a> to know more about them! <hr /> <h2 id="finetuning">Fine-tuning</h2> We'll use <a href="https://huggingface.co/datasets/nielsr/breast-cancer">this dataset</a> consisting of breast cancer scans. In the medical imaging domain, being able to segment the cells containing malignancy is an important task. <h3 id="data-preparation">Data preparation</h3> Let's first get the dataset. <div class="codehilite"><pre><code>remote_path = "https://hf.co/datasets/sayakpaul/sample-datasets/resolve/main/breast-cancer-dataset.tar.gz" dataset_path = keras.utils.get_file( "breast-cancer-dataset.tar.gz", remote_path, untar=True ) </code></pre></div> Let's now visualize a sample from the dataset. (The <code>show_mask()</code> utility is taken from <a href="https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SAM/Fine_tune_SAM_(segment_anything)_on_a_custom_dataset.ipynb">this notebook</a>) <div class="codehilite"><pre><code>def show_mask(mask, ax, random_color=False): if random_color: color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) else: color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6]) h, w = mask.shape[-2:] mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) ax.imshow(mask_image) # Load all the image and label paths. image_paths = sorted(glob.glob(os.path.join(dataset_path, "images/*.png"))) label_paths = sorted(glob.glob(os.path.join(dataset_path, "labels/*.png"))) # Load the image and label. idx = 15 image = Image.open(image_paths[idx]) label = Image.open(label_paths[idx]) image = np.array(image) ground_truth_seg = np.array(label) # Display. fig, axes = plt.subplots() axes.imshow(image) show_mask(ground_truth_seg, axes) axes.title.set_text(f"Ground truth mask") axes.axis("off") plt.show() tf.shape(ground_truth_seg) </code></pre></div> <img alt="png" src="/img/examples/vision/sam/sam_26_0.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code><tf.Tensor: shape=(2,), dtype=int32, numpy=array([256, 256], dtype=int32)> </code></pre></div> </div> <h3 id="tfdatadataset">Preparing <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset"><code>tf.data.Dataset</code></a></h3> We now write a generator class to prepare the images and the segmentation masks using the <code>processor</code> utilized above. We will leverage this generator class to create a <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset"><code>tf.data.Dataset</code></a> object for our training set by using <code>tf.data.Dataset.from_generator()</code>. Utilities of this class have been adapted from <a href="https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SAM/Fine_tune_SAM_(segment_anything)_on_a_custom_dataset.ipynb">this notebook</a>. The generator is responsible for yielding the preprocessed images and the segmentation masks, and some other metadata needed by the SAM model. <div class="codehilite"><pre><code>class Generator: """Generator class for processing the images and the masks for SAM fine-tuning.""" def __init__(self, dataset_path, processor): self.dataset_path = dataset_path self.image_paths = sorted( glob.glob(os.path.join(self.dataset_path, "images/*.png")) ) self.label_paths = sorted( glob.glob(os.path.join(self.dataset_path, "labels/*.png")) ) self.processor = processor def __call__(self): for image_path, label_path in zip(self.image_paths, self.label_paths): image = np.array(Image.open(image_path)) ground_truth_mask = np.array(Image.open(label_path)) # get bounding box prompt prompt = self.get_bounding_box(ground_truth_mask) # prepare image and prompt for the model inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="np") # remove batch dimension which the processor adds by default inputs = {k: v.squeeze(0) for k, v in inputs.items()} # add ground truth segmentation inputs["ground_truth_mask"] = ground_truth_mask yield inputs def get_bounding_box(self, ground_truth_map): # get bounding box from mask y_indices, x_indices = np.where(ground_truth_map > 0) x_min, x_max = np.min(x_indices), np.max(x_indices) y_min, y_max = np.min(y_indices), np.max(y_indices) # add perturbation to bounding box coordinates H, W = ground_truth_map.shape x_min = max(0, x_min - np.random.randint(0, 20)) x_max = min(W, x_max + np.random.randint(0, 20)) y_min = max(0, y_min - np.random.randint(0, 20)) y_max = min(H, y_max + np.random.randint(0, 20)) bbox = [x_min, y_min, x_max, y_max] return bbox </code></pre></div> <code>get_bounding_box()</code> is responsible for turning the ground-truth segmentation maps into bounding boxes. These bounding boxes are fed to SAM as prompts (along with the original images) during fine-tuning and SAM is then trained to predict valid masks. The advantage of this first creating a generator and then using it to create a <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset"><code>tf.data.Dataset</code></a> is the flexbility. Sometimes, we may need to use utitlities from other libraries (<a href="https://albumentations.ai/"><code>albumentations</code></a>, for example) which may not come in native TensorFlow implementations. By using this workflow, we can easily accommodate such use case. But the non-TF counterparts might introduce performance bottlenecks, though. However, for our example, it should work just fine. Now, we prepare the <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset"><code>tf.data.Dataset</code></a> from our training set. <div class="codehilite"><pre><code># Define the output signature of the generator class. output_signature = { "pixel_values": tf.TensorSpec(shape=(3, None, None), dtype=tf.float32), "original_sizes": tf.TensorSpec(shape=(None,), dtype=tf.int64), "reshaped_input_sizes": tf.TensorSpec(shape=(None,), dtype=tf.int64), "input_boxes": tf.TensorSpec(shape=(None, None), dtype=tf.float64), "ground_truth_mask": tf.TensorSpec(shape=(None, None), dtype=tf.int32), } # Prepare the dataset object. train_dataset_gen = Generator(dataset_path, processor) train_ds = tf.data.Dataset.from_generator( train_dataset_gen, output_signature=output_signature ) </code></pre></div> Next, we configure the dataset for performance. <div class="codehilite"><pre><code>auto = tf.data.AUTOTUNE batch_size = 2 shuffle_buffer = 4 train_ds = ( train_ds.cache() .shuffle(shuffle_buffer) .batch(batch_size) .prefetch(buffer_size=auto) ) </code></pre></div> Take a single batch of data and inspect the shapes of the elements present inside of it. <div class="codehilite"><pre><code>sample = next(iter(train_ds)) for k in sample: print(k, sample[k].shape, sample[k].dtype, isinstance(sample[k], tf.Tensor)) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>pixel_values (2, 3, 1024, 1024) <dtype: 'float32'> True original_sizes (2, 2) <dtype: 'int64'> True reshaped_input_sizes (2, 2) <dtype: 'int64'> True input_boxes (2, 1, 4) <dtype: 'float64'> True ground_truth_mask (2, 256, 256) <dtype: 'int32'> True </code></pre></div> </div> <h3 id="training">Training</h3> We will now write DICE loss. This implementation is based on <a href="https://docs.monai.io/en/stable/losses.html#diceloss">MONAI DICE loss</a>. <div class="codehilite"><pre><code>def dice_loss(y_true, y_pred, smooth=1e-5): y_pred = tf.sigmoid(y_pred) reduce_axis = list(range(2, len(y_pred.shape))) if batch_size > 1: # reducing spatial dimensions and batch reduce_axis = [0] + reduce_axis intersection = tf.reduce_sum(y_true * y_pred, axis=reduce_axis) y_true_sq = tf.math.pow(y_true, 2) y_pred_sq = tf.math.pow(y_pred, 2) ground_o = tf.reduce_sum(y_true_sq, axis=reduce_axis) pred_o = tf.reduce_sum(y_pred_sq, axis=reduce_axis) denominator = ground_o + pred_o # calculate DICE coefficient loss = 1.0 - (2.0 * intersection + 1e-5) / (denominator + 1e-5) loss = tf.reduce_mean(loss) return loss </code></pre></div> <h2 id="finetuning-sam">Fine-tuning SAM</h2> We will now fine-tune SAM's decoder part. We will freeze the vision encoder and prompt encoder layers. <div class="codehilite"><pre><code># initialize SAM model and optimizer sam = TFSamModel.from_pretrained("facebook/sam-vit-base") optimizer = keras.optimizers.Adam(1e-5) for layer in sam.layers: if layer.name in ["vision_encoder", "prompt_encoder"]: layer.trainable = False @tf.function def train_step(inputs): with tf.GradientTape() as tape: # pass inputs to SAM model outputs = sam( pixel_values=inputs["pixel_values"], input_boxes=inputs["input_boxes"], multimask_output=False, training=True, ) predicted_masks = tf.squeeze(outputs.pred_masks, 1) ground_truth_masks = tf.cast(inputs["ground_truth_mask"], tf.float32) # calculate loss over predicted and ground truth masks loss = dice_loss(tf.expand_dims(ground_truth_masks, 1), predicted_masks) # update trainable variables trainable_vars = sam.trainable_variables grads = tape.gradient(loss, trainable_vars) optimizer.apply_gradients(zip(grads, trainable_vars)) return loss </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>All model checkpoint layers were used when initializing TFSamModel. </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>All the layers of TFSamModel were initialized from the model checkpoint at facebook/sam-vit-base. If your task is similar to the task the model of the checkpoint was trained on, you can already use TFSamModel for predictions without further training. WARNING:absl:At this time, the v2.11+ optimizer [`tf.keras.optimizers.Adam`](/api/optimizers/adam#adam-class) runs slowly on M1/M2 Macs, please use the legacy Keras optimizer instead, located at [`tf.keras.optimizers.legacy.Adam`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/legacy/Adam). </code></pre></div> </div> We can now run the training for three epochs. We might have a warning about gradients not existing on IoU prediction head of mask decoder, we can safely ignore that. <div class="codehilite"><pre><code># run training for epoch in range(3): for inputs in train_ds: loss = train_step(inputs) print(f"Epoch {epoch + 1}: Loss = {loss}") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? WARNING:tensorflow:Gradients do not exist for variables ['tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_in/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/proj_out/bias:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/kernel:0', 'tf_sam_model_1/mask_decoder/iou_prediction_head/layers_._0/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument? Epoch 1: Loss = 0.08322787284851074 Epoch 2: Loss = 0.05677264928817749 Epoch 3: Loss = 0.07764029502868652 </code></pre></div> </div> <h3 id="serialize-the-model">Serialize the model</h3> We serialized the model and pushed for you below. <code>push_to_hub</code> method serializes model, generates a model card and pushes it to Hugging Face Hub, so that other people can load the model using <code>from_pretrained</code> method to infer or further fine-tune. We also need to push the same preprocessor in the repository. Find the model and the preprocessor <a href="https://huggingface.co/merve/sam-finetuned">here</a>. <div class="codehilite"><pre><code># sam.push_to_hub("merve/sam-finetuned") # processor.push_to_hub("merve/sam-finetuned") </code></pre></div> We can now infer with the model. <div class="codehilite"><pre><code># Load another image for inference. idx = 20 raw_image_inference = Image.open(image_paths[idx]) # process the image and infer preprocessed_img = processor(raw_image_inference) outputs = sam(preprocessed_img) </code></pre></div> Lastly, we can visualize the results. <div class="codehilite"><pre><code>infer_masks = outputs["pred_masks"] iou_scores = outputs["iou_scores"] show_masks_on_image(raw_image_inference, masks=infer_masks, scores=iou_scores) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). WARNING:matplotlib.image:Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers). </code></pre></div> </div> <img alt="png" src="/img/examples/vision/sam/sam_48_1.png" /> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#segment-anything-model-with-🤗transformers'>Segment Anything Model with 🤗Transformers</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#installation'>Installation</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#sam-in-a-few-words'>SAM in a few words</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#running-inference-with-sam'>Running inference with SAM</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#finetuning'>Fine-tuning</a> </div> <div class='k-outline-depth-3'> <a href='#data-preparation'>Data preparation</a> </div> <div class='k-outline-depth-3'> <a href='#preparing-tfdatadataset'>Preparing <code>tf.data.Dataset</code></a> </div> <div class='k-outline-depth-3'> <a href='#training'>Training</a> </div> <div class='k-outline-depth-3'> <a href='#serialize-the-model'>Serialize the model</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>