Interactive Segmentation

Module: Computer Vision | Difficulty: Intermediate

Segment Anything Model (SAM)

Prompt Encodings

• Point prompt:
• Box prompt:
• Text prompt: CLIP embedding

Mask Decoder

IoU Prediction

Click Accumulation

import torch
import torch.nn as nn

class SimpleSAM(nn.Module):
    def __init__(self, image_encoder, prompt_encoder, mask_decoder):
        super().__init__()
        self.image_encoder = image_encoder
        self.prompt_encoder = prompt_encoder
        self.mask_decoder = mask_decoder
    
    def forward(self, image, points=None, boxes=None):
        image_feat = self.image_encoder(image)
        
        if points is not None:
            prompt_feat = self.prompt_encoder(points)
        elif boxes is not None:
            prompt_feat = self.prompt_encoder(boxes)
        
        mask, iou = self.mask_decoder(image_feat, prompt_feat)
        return mask, iou

Key Takeaways

• SAM provides universal segmentation with simple prompts
• Iterative clicking refines the segmentation mask
• SAM generalizes to unseen objects without fine-tuning