Lesson 53: Segment Anything¶
Lesson 43's instance segmentation trained a network to output every instance's mask at once, then untangled them with center-voting. Segment Anything (SAM) (Kirillov et al., 2023) takes a different approach: segment on demand. Give the model a prompt — a point, a box, a rough scribble — and it returns the mask of whichever single object that prompt points to. This lesson builds SAM's core idea: inject the prompt as an explicit spatial signal alongside the image, and train the network to produce exactly the prompted object's mask, nothing else.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Scenes with a point prompt¶
Reuse Lesson 43's overlapping-circles scenes. Instead of predicting all instances at once, sample one point inside one of the instances (the "prompt") and train the model to output only that instance's mask.
SIZE = 32
def make_scene(rng, size=SIZE, n_circles=2, r=5):
scene = np.zeros((size, size), dtype=np.float32)
inst_mask = np.zeros((size, size), dtype=np.int64)
centers = []
for k in range(n_circles):
if k == 0:
cx, cy = rng.integers(r + 8, size - r - 8), rng.integers(r + 8, size - r - 8)
else:
px, py = centers[-1]
angle = rng.uniform(0, 2 * np.pi)
dist = rng.uniform(6, 9) # overlapping but centers still separable
cx = int(np.clip(px + dist * np.cos(angle), r, size - r - 1))
cy = int(np.clip(py + dist * np.sin(angle), r, size - r - 1))
centers.append((cx, cy))
yy, xx = np.mgrid[0:size, 0:size]
m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2
scene[m] = 1.0
inst_mask[m] = k + 1
scene = np.clip(scene + rng.normal(0, 0.05, scene.shape), 0, 1).astype(np.float32)
return scene, inst_mask, centers
def sample_point_prompt(inst_mask, instance_id, rng):
ys, xs = np.where(inst_mask == instance_id)
i = rng.integers(len(xs))
return xs[i], ys[i]
rng = np.random.default_rng(17)
N = 400
scenes, inst_masks, all_centers, prompts, prompt_targets = [], [], [], [], []
for _ in range(N):
s, im, c = make_scene(rng)
scenes.append(s); inst_masks.append(im); all_centers.append(c)
inst_id = rng.integers(1, len(c) + 1) # pick a random instance to prompt for
px, py = sample_point_prompt(im, inst_id, rng)
prompts.append((px, py))
prompt_targets.append((im == inst_id).astype(np.float32))
scenes = np.array(scenes, dtype=np.float32)
prompt_targets = np.array(prompt_targets, dtype=np.float32)
prompts_arr = np.array(prompts, dtype=np.float32)
split = int(0.85 * N)
Xtr, Ptr, Mtr = scenes[:split], prompts_arr[:split], prompt_targets[:split]
Xte, Pte, Mte = scenes[split:], prompts_arr[split:], prompt_targets[split:]
inst_te, centers_te = inst_masks[split:], all_centers[split:]
fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for i in range(4):
axes[i].imshow(scenes[i], cmap='gray')
axes[i].scatter([prompts[i][0]], [prompts[i][1]], c='red', marker='*', s=100)
axes[i].axis('off')
plt.suptitle('Scenes with their point prompt (red star)')
plt.show()
Injecting the prompt as a spatial signal¶
SAM's architecture is three pieces: an image encoder (a heavy, expensive-to-run backbone, computed once per image), a lightweight prompt encoder (turns a point/box/mask into an embedding), and a fast mask decoder that combines the two. Real SAM's image encoder is always a Vision Transformer (Lesson 47's ViT, MAE-pretrained in the original paper) — never a CNN — so this lesson's encoder is one too: patchify the image exactly as in Lesson 47, run a small Transformer encoder over the patch sequence, then reshape the resulting per-patch tokens back into a spatial grid of features. The one design choice that matters most for a point prompt: encode where the point is spatially, not as an abstract vector broadcast uniformly over the whole image. Rendering the prompt as a Gaussian blob on its own channel, patchified together with the image, gives every patch embedding — and therefore every attention computation downstream — a direct, position-aware cue about the prompt from the very first layer.
def point_prompt_map(points, size=SIZE):
B = points.shape[0]
yy, xx = torch.meshgrid(torch.arange(size), torch.arange(size), indexing='ij')
yy = yy.float().unsqueeze(0).expand(B, -1, -1)
xx = xx.float().unsqueeze(0).expand(B, -1, -1)
px, py = points[:, 0].view(B, 1, 1), points[:, 1].view(B, 1, 1)
d2 = (xx - px) ** 2 + (yy - py) ** 2
return torch.exp(-d2 / (2 * 2.0 ** 2)) # a Gaussian blob centered at the prompt point
def patchify(img, patch_size):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, P, P)
patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)
return patches.reshape(B, -1, C * P * P)
def positional_encoding(T, D):
pos = torch.arange(T).unsqueeze(1).float()
i = torch.arange(D).unsqueeze(0).float()
angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))
angles = pos * angle_rates
pe = torch.zeros(T, D)
pe[:, 0::2] = torch.sin(angles[:, 0::2])
pe[:, 1::2] = torch.cos(angles[:, 1::2])
return pe
class TinySAM(nn.Module):
def __init__(self, img_size=SIZE, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.patch_size = patch_size
self.grid = img_size // patch_size
patch_dim = 2 * patch_size * patch_size # image channel + prompt channel, patchified together
self.embed = nn.Linear(patch_dim, embed_dim)
self.register_buffer('pos_enc', positional_encoding(self.grid * self.grid, embed_dim))
layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,
batch_first=True, dropout=0.0)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.mask_decoder = nn.Sequential(
nn.Conv2d(embed_dim, 16, 3, padding=1), nn.ReLU(),
nn.Conv2d(16, 1, 1),
)
def forward(self, img, point):
prompt_map = point_prompt_map(point).unsqueeze(1)
x = torch.cat([img, prompt_map], dim=1)
tok = self.embed(patchify(x, self.patch_size)) + self.pos_enc
out = self.encoder(tok) # (B, n_patches, embed_dim)
feat = out.transpose(1, 2).reshape(x.shape[0], -1, self.grid, self.grid)
feat_up = F.interpolate(feat, size=img.shape[-2:], mode='nearest') # back to full resolution
return self.mask_decoder(feat_up).squeeze(1)
torch.manual_seed(0)
model = TinySAM()
opt = torch.optim.Adam(model.parameters(), lr=0.003)
Xt = torch.tensor(Xtr).unsqueeze(1); Pt = torch.tensor(Ptr); Mt = torch.tensor(Mtr)
for _ in range(600):
opt.zero_grad()
loss = F.binary_cross_entropy_with_logits(model(Xt, Pt), Mt)
loss.backward()
opt.step()
with torch.no_grad():
logits_te = model(torch.tensor(Xte).unsqueeze(1), torch.tensor(Pte))
preds = (logits_te > 0).numpy().astype(bool)
def iou(a, b):
inter, union = (a & b).sum(), (a | b).sum()
return inter / union if union > 0 else float('nan')
ious = [iou(preds[i], Mte[i].astype(bool)) for i in range(len(Xte))]
print(f'point-prompted mask IoU: mean = {np.mean(ious):.3f}')
The real test: does the prompt actually select an instance?¶
Take one test scene with two overlapping circles, and query it twice with two different point prompts — one inside each circle. A genuinely prompt-driven model should return two different masks; a model that's secretly ignoring the prompt and just doing semantic segmentation (Lesson 42) would return the same "all circle pixels" mask both times.
idx = 0
scene0 = Xte[idx]
(c1x, c1y), (c2x, c2y) = centers_te[idx]
with torch.no_grad():
logit_a = model(torch.tensor(scene0[None, None]), torch.tensor([[c1x, c1y]], dtype=torch.float32))
logit_b = model(torch.tensor(scene0[None, None]), torch.tensor([[c2x, c2y]], dtype=torch.float32))
mask_a, mask_b = (logit_a[0] > 0).numpy(), (logit_b[0] > 0).numpy()
overlap = (mask_a & mask_b).sum() / max((mask_a | mask_b).sum(), 1)
print(f'mask A size: {mask_a.sum()} px, mask B size: {mask_b.sum()} px')
print(f'overlap (IoU) between the two prompted masks: {overlap:.3f} (low = genuinely different instances)')
fig, axes = plt.subplots(1, 3, figsize=(9, 3))
axes[0].imshow(scene0, cmap='gray')
axes[0].scatter([c1x, c2x], [c1y, c2y], c=['red', 'cyan'], marker='*', s=100)
axes[0].set_title('scene, both prompts'); axes[0].axis('off')
axes[1].imshow(mask_a, cmap='gray')
axes[1].scatter([c1x], [c1y], c='red', marker='*', s=100)
axes[1].set_title('mask for red prompt'); axes[1].axis('off')
axes[2].imshow(mask_b, cmap='gray')
axes[2].scatter([c2x], [c2y], c='cyan', marker='*', s=100)
axes[2].set_title('mask for cyan prompt'); axes[2].axis('off')
plt.show()
The two prompts on the same, unchanged image produce two clearly different masks, each tightly following the circle the prompt actually landed in — the model is genuinely conditioning on the prompt, not just re-running semantic segmentation and ignoring it.
This reframes instance segmentation entirely: Lesson 43 needed a specialized clustering step (vote, then assign) to recover instances from a single forward pass over the whole image. SAM sidesteps that by making "which instance" an input rather than something to be inferred — ambiguity about what a click means is resolved by the human (or an upstream detector, Lesson 41) providing the point, not by the segmentation model guessing. This is also what makes SAM's "segment anything" claim viable: it was never trained on specific object categories at all, only on the general task "here's a point, here's its object's boundary" across a huge, diverse dataset (SA-1B, over 1 billion masks) — a category-agnostic skill that transfers to object types never seen during training, much like Lesson 55's zero-shot classification transferred to unseen class names.
Exercises¶
- Change the prompt sampling in the training loop to always pick a point near an instance's edge rather than anywhere inside it (
sample_point_promptcurrently samples uniformly). Does edge-prompt IoU come out lower than interior-prompt IoU, and why would a network find edge points more ambiguous? - SAM in practice supports box prompts as well as point prompts, and typically outputs several candidate masks per prompt (since a single point can be ambiguous — a point on a shirt could mean "the shirt" or "the whole person"). Sketch how you'd change
TinySAM's final layer to output 3 candidate masks instead of 1, and what additional training signal you'd need to know which of the 3 to prefer. - Increase
n_circlesfrom 2 to 4 inmake_scene, keeping everything else the same, and rerun the two-different-prompts test on a 4-circle scene. Does the model still cleanly separate two specific circles by their prompts, or does more clutter degrade prompt selectivity?