TDM 19000: Introduction to SAM

The Segment Anything Model (SAM) is a AI foundational model that is designed for image segmentation. This model is able to identify and isolate any object within an image or video with zero-shot generalization.

Zero-shot generalization is where an AI model correctlly performs a task or classifies data into categories that it hadn’t encountered during any of its training. It’s able to understand the relationship between data it has seen and data it has never seen.

Previously, we were using YOLO to define boxes for the objects. We can use SAM take the bounding boxes, find which pixels belong to each, and then create masks for the defined objects. This method of pixel-level segmentation allows for more precision when processing an image.

To start, load in Meta’s SAM model:

import torch
from transformers import SamModel, SamProcessor

device = "cuda" if torch.cuda.is_available() else "cpu"

sam_model = SamModel.from_pretrained("facebook/sam-vit-base")
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
sam_model.to(device)

We’ve got to save the original image, and then put it into the SAM processor along with the bounding boxes (saved as input_boxes).

original_img = Image.open(image_path)

sam_inputs = sam_processor([image name], input_boxes=[[your input boxes]], return_tensors="pt")

In the SAM processor, make sure to keep the name of your input boxes within a set of []. Otherwise this will not run correctly.

For example, input_boxes=[input_boxes].

When running the model: SAM will predict masks for each box.

with torch.no_grad():
    outputs = sam_model(**sam_inputs)

Once the masks exist, they will need to be post-processed. In this case, this means rescaling the masks to the size that makes sense with the actual image, and ensuring that there is one mask per input box. To do this:

masks = sam_processor.image_processor.post_process_masks(
    outputs.pred_masks,
    sam_inputs["original_sizes"],
    sam_inputs["reshaped_input_sizes"]
)[0]

These masks should be closer to the actual shape of the objects being predicted in the image. To visualize them onto the original image, we have to:

  • Display the image (as the background layer)

  • Fix mask shape (mask should be a 2D array)

  • Overlay on image

m = masks[0]

plt.imshow(original_img)

if m.ndim == 3:
    m = m[0]

plt.imshow(m, alpha=0.4)
plt.axis("off")
plt.show()

This will display the mask for one selected object in the image. You can change m = masks[0] to other values, for instance, m = masks[1] or m = masks[12] to see the masks for other selected objects in the image.