Lesson 51: Segment Anything¶
Lesson 42'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 42'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. 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, concatenated with the image, gives every convolution a direct, position-aware cue.
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
class TinySAM(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(2, 16, 3, padding=1), nn.ReLU(), # image channel + prompt channel
nn.Conv2d(16, 16, 3, padding=1), nn.ReLU(),
nn.Conv2d(16, 1, 1),
)
def forward(self, img, point):
prompt_map = point_prompt_map(point).unsqueeze(1)
return self.net(torch.cat([img, prompt_map], dim=1)).squeeze(1)
torch.manual_seed(0)
model = TinySAM()
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xt = torch.tensor(Xtr).unsqueeze(1); Pt = torch.tensor(Ptr); Mt = torch.tensor(Mtr)
for _ in range(300):
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 41) 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 42 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 40) 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 49's zero-shot classification transferred to unseen class names.
Exercise¶
- 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?