Say Goodbye to Tedious Annotation: Annotate Object for Segmentation By Words

Have you ever spent hours annotating objects in images for segmentation tasks? It’s a tedious and time-consuming process that can drain your productivity and creativity. But what if you could annotate objects for segmentation with just a word? Imagine being able to create high-quality segmentation data in minutes instead of hours, with no manual annotation required. Thanks to recent advances in AI, this is now possible. In this blog post, we’ll introduce you to our groundbreaking AI pipeline that lets you create semantic segmentation data for objects using only words. With our pipeline, you can say goodbye to tedious annotation and hello to effortless object segmentation. Let’s dive in!

Grounding DINO

Grounding DINO is a neural network architecture that combines a transformer-based language model (similar to GPT) with a computer vision model for object detection. The goal of grounding DINO is to enable zero-shot object detection, which means that the model can detect objects that it hasn’t been specifically trained to recognize.

The way grounding DINO works is by first encoding a natural language phrase that describes an object, such as “a red car,” into a vector representation using the transformer-based language model. Then, the vector is passed through a computer vision model (typically a convolutional neural network) that processes an input image and produces a feature map. The vector is “grounded” to the feature map by finding the spatial location that corresponds to the object described by the phrase.

The result is a bounding box that indicates the location of the object in the image. Because the language model is trained on a large corpus of text, grounding DINO can detect a wide range of objects without the need for extensive training data.

Segment Anything Model

The “Segment Anything” model is a computer vision model that uses a transformer-based architecture with self-attention mechanisms to segment objects, their attributes, and their relationships within an image. The model includes an encoder-decoder architecture, where the encoder applies a self-attention mechanism to capture the relationships between objects and the decoder generates the segmentation output using the encoded features. The model also has the option to provide the bounding box coordinates or polygonal coordinates to perform segmentation only for a specific area.

Methodology

The goal of this method is to create a semantic segmentation dataset for a particular object in images. We achieve this by using two models: Grounding DINO and Segment Anything.

Grounding DINO is a zero-shot object detection model that can detect objects in images based on textual input. It takes an image and a text description of the object to be detected as input and returns the bounding boxes for each instance of the object in the image. For example, if we want to generate a segmentation dataset for apples, we can pass an image containing apples and the text “apple” to the Grounding DINO model, and it will return the bounding boxes for each instance of an apple in the image.

Once we have the bounding boxes for each instance of the object in the image, we can use the Segment Anything model to generate binary masks for each instance. Segment Anything is a segmentation model recently open-sourced by Meta that can accept a bounding box as input and generate a binary mask for the corresponding object.

To generate the binary masks, we first initialize a mask image with pixel values zero everywhere. We then iterate over the bounding boxes given by the Grounding DINO model and generate a temporary mask of the object instance using these boxes. We copy the values from the temporary mask to the mask we initialized earlier. Since the size of the temporary mask is the same as the mask initialized, copying each instance of an object to the mask is equivalent to generating the masked area for all objects in one go.

After we have processed all instances of the object in the image, we save the final mask as the annotation for the image. The mask will have pixel values of zero everywhere except the regions in which the object is present.

Using this methodology, we can generate semantic segmentation datasets for any object by passing the object’s name to the Grounding DINO model and using the Segment Anything model to generate the binary masks for each instance of the object in the image. This approach offers a simple and effective way to generate high-quality segmentation datasets without the need for extensive manual annotation.

Implementation in Code

Install Segment Anything Model (SAM) and its dependencies

!pip install 'git+https://github.com/facebookresearch/segment-anything.git'
!pip install -q dataclasses-json supervision

Here we are installing the segment anything python package open sourced by the Meta which provides the util functions for loading the SAM model and making predictions with it. The second command installs other Python packages dataclasses-json, and supervision which are required for visualization purpose.

Clone and Install the Dependencies for Grounding Dino

          !git clone https://github.com/IDEA-Research/GroundingDINO.git
%cd GroundingDINO

# we use latest Grounding DINO model API that is not official yet
!git checkout feature/more_compact_inference_api

!pip install -q -e .
!pip install -q dataclasses-json onemetric     

Here we are installing the grounding DINO requirements to load and run the Grounding DINO model.

Download the Pretrained Weights of the Models

          #Download Segment Anything Model weights.
!wget -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

#Download Grounding DINO weights.
!wget -q https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth

Importing Libraries

          import torch
import os
import supervision
import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
from groundingdino.util.inference import Model
%matplotlib inline    

Here we are importing the required libraries. Here we have imported libraries like opencv, torch,numpy, and os for loading and processing images and tensors. We also import the respective classes required by both models. The supervision library provides a lot of options to plot and visualize masks on our images.

Loading Both Models

          #Check if GPU is available.
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

#Load the SAM model.
CHECKPOINT_PATH_SAM ='/content/GroundingDINO/sam_vit_h_4b8939.pth'
MODEL_TYPE = "vit_h"
sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH_SAM).to(device=DEVICE)
mask_predictor = SamPredictor(sam)

#Load the Grounding DINO model.
WEIGHTS_PATH_GDINNO = "/content/GroundingDINO/groundingdino_swint_ogc.pth"
CONFIG_PATH = "/content/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py"
model = Model(model_config_path=CONFIG_PATH, model_checkpoint_path=WEIGHTS_PATH_GDINNO)    

Here we first check if any GPU device is available then we load both SAM and Grounding DINO model. In case GPU is available we transfer the model to GPU for faster execution.

Performing Object Detection With Grounding DINO

          #Source image path and other hyperparameters.
SOURCE_IMAGE_PATH = f"/content/apple.jpg"
CLASSES = ['apple']
BOX_TRESHOLD = 0.35
TEXT_TRESHOLD = 0.25


#Performing detection with Grounding DINO.
image = cv2.imread(SOURCE_IMAGE_PATH)
caption = ", ".join(CLASSES)

detections_dino, labels = model.predict_with_caption(
    image=image,
    caption=caption,
    box_threshold=BOX_TRESHOLD,
    text_threshold=TEXT_TRESHOLD
)

Here first we start off by setting the path to the source image and the name of the classes. Since we just want to detect apple we are passing text apple here. With grounding DINO you can combine different classes of information with words. Say we want to detect a man sitting on a chair then we can simply pass this prompt to the model and the model will detect the same.

Here the model will return bounding box coordinates and the labels. We would use these bounding boxes to generate the segmentation for all instances of apple present in the image.

Generate masks with SAM

For generating the masks for the object I tried passing all the bounding boxes at once to the SAM mask predictor but it threw some exceptions hence I proceeded with predicting one bounding box at a time. Please let me know in the comment section if anyone is able to process more than one box at a time.

          #Annotate for all occurances.

image_bgr = cv2.imread('/content/apple.jpg')
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
seg_mask=np.zeros_like(image_rgb)
pred_mask_list=[]

for i in range(detections_dino.xyxy.shape[0]):
  box=detections_dino.xyxy[i]
  mask_predictor.set_image(image_rgb)

  masks, scores, logits = mask_predictor.predict(
      box=box,
      multimask_output=True
  )

  masks=np.transpose(masks,[2,1,0])
  masks=np.transpose(masks,[1,0,2])

  pred_mask_list.append(masks)
        

Here we iterate over the detections_dino object which stores bounding boxes in xyxy attribute. Before doing that we initialize an empty mask seg_mask having pixel values zero all over except in the region where apple would exist in the image having a value of 1 or 255 which would act as the segmentation label.

In the iteration part, we would predict the mask for each bounding box at a time and append it to a list called pred_mask_list.

In the above figure, the top left image is the input image for which the segmentation mask needs to be generated and the top right image is the final mask obtained after the whole process. But since we can process one bounding box at a time we would generate the segmentation mask separately for them and append it to the list which is shown in the bottom right and left in the above image.

Copying the Individual Mask to the Main Mask

We saw in the above section that we generate the mask for all instances of the object separately hence these masks must be combined in order to get the final mask which contains the masks for all instances of the input class.

          #Setting all the masks places to 255.

for i in range(len(pred_mask_list)):
  seg_mask[pred_mask_list[i]]=255 

We iterate over each mask and using indexing we copy all 255 pixels to the main mask seg_mask which results in forming the final mask which is shown in the top right corner of the above figure.

Further Improvements

* Currently, we see here that we are generating data for a single class and we are only generating data for semantic segmentation. Hence the same pipeline with little logic modification can be extended to multiple classes and for instance segmentation.

* Currently, we are creating a segmentation mask for one bounding box at a time which is a little inefficient hence we can edit the prediction code in SAM code to process multiple boxes at a time.

References and Further Readings

How to Use the Segment Anything Model (SAM)

Segment Anything (SAM) is a computer vision model developed by Meta AI. In this guide, you will learn how to use SAM on your own data.

blog.roboflow.com

GitHub - IDEA-Research/GroundingDINO: [ECCV 2024] Official implementation of the paper "Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection"

[ECCV 2024] Official implementation of the paper "Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection" - IDEA-Research/GroundingDINO

github.com

Link to Notebook:

GitHub - VikasOjha666/Automated_Segmentation_Annotation

Contribute to VikasOjha666/Automated_Segmentation_Annotation development by creating an account on GitHub.

github.com

Congrats you have just learned an amazing concept with which you can speed up your segmentation project and minimize your effort. If you have any questions regarding this post please feel free to leave in the comments section I will try to answer the same.