Lesson 54: Promptable Video Segmentation¶
Lesson 53's Segment Anything Model turns a point or box prompt into a mask for a single image. This lesson extends that same promptable idea to video, in the style of SAM 2: prompt an object once, on any frame, and propagate its mask forward (and backward) through the rest of the video by tracking it from frame to frame. The key new mechanism is a streaming memory — a compact representation of the object's appearance from previous frames that the model attends to when segmenting the current one, so the object can be re-found after occlusion or appearance change without being re-prompted every frame.
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
A synthetic video with occlusion¶
A small circle (the target) moves in a straight line across a 32x32 clip, passing directly behind a static gray square (a stand-in for anything that happens to occlude the camera's view, like another object or scene clutter) for a few consecutive frames before re-emerging on the far side. The ground-truth mask always reflects the target's true position, even while it's invisible in the pixels — exactly the situation SAM 2 is built for: the object is still there, the video just doesn't show it for a while.
SIZE = 32
T = 14
R = 4
OCC_CENTER = np.array([SIZE / 2.0, SIZE / 2.0])
OCC_HALF = 3.0 # the static occluder is a (2*OCC_HALF)-wide square at the center of the frame
def in_occluder(p):
return abs(p[0] - OCC_CENTER[0]) <= OCC_HALF and abs(p[1] - OCC_CENTER[1]) <= OCC_HALF
def onscreen(p, r=R, size=SIZE):
return r <= p[0] <= size - r and r <= p[1] <= size - r
def render_frame(target_pos, visible, size=SIZE, r=R):
img = np.zeros((size, size), dtype=np.float32)
mask = np.zeros((size, size), dtype=np.float32)
yy, xx = np.mgrid[0:size, 0:size]
target_region = ((xx - target_pos[0]) ** 2 + (yy - target_pos[1]) ** 2) <= r ** 2
mask[target_region] = 1.0 # the mask always reflects the TRUE position, visible or not
if visible:
img[target_region] = 1.0
occ_region = (np.abs(xx - OCC_CENTER[0]) <= OCC_HALF) & (np.abs(yy - OCC_CENTER[1]) <= OCC_HALF)
img[occ_region] = 0.55 # the occluder is drawn on top, regardless of the target
return img, mask
def make_clip(rng, size=SIZE, r=R, T=T):
"""A target crosses straight through the occluder, becoming briefly invisible partway
through the clip, then re-emerges and keeps moving in the same direction."""
for _attempt in range(300):
edge = rng.integers(0, 4)
if edge == 0:
p0 = np.array([rng.uniform(r + 1, size - r - 1), r + 1.0])
elif edge == 1:
p0 = np.array([rng.uniform(r + 1, size - r - 1), size - r - 1.0])
elif edge == 2:
p0 = np.array([r + 1.0, rng.uniform(r + 1, size - r - 1)])
else:
p0 = np.array([size - r - 1.0, rng.uniform(r + 1, size - r - 1)])
jitter = rng.uniform(-2, 2, size=2)
direction = (OCC_CENTER + jitter) - p0
direction = direction / np.linalg.norm(direction)
speed = rng.uniform(2.2, 2.8)
traj = [p0 + direction * speed * t for t in range(T)]
visible = [not in_occluder(p) for p in traj]
occ_frames = [t for t, v in enumerate(visible) if not v]
if not occ_frames or not (2 <= len(occ_frames) <= 3):
continue
if occ_frames[0] < 2 or occ_frames[-1] > T - 4:
continue
if not all(onscreen(p, r, size) for p in traj):
continue
frames, masks = [], []
for t in range(T):
img, mask = render_frame(traj[t], visible[t], size, r)
frames.append(img)
masks.append(mask)
return np.array(frames, dtype=np.float32), np.array(masks, dtype=np.float32), occ_frames, traj
raise RuntimeError('failed to generate a clean clip')
example_frames, example_masks, example_occ, example_traj = make_clip(np.random.default_rng(0))
fig, axes = plt.subplots(1, T, figsize=(16, 1.6))
for t, ax in enumerate(axes):
ax.imshow(example_frames[t], cmap='gray', vmin=0, vmax=1)
ax.axis('off')
ax.set_title('occ' if t in example_occ else '', fontsize=8, color='red')
plt.suptitle('One clip: the target crosses behind the static occluder (frames marked "occ")', y=1.15)
plt.show()
Prompting once, and remembering what was prompted¶
Exactly as in Lesson 53, the user prompts only once — here, the ground-truth mask on frame 0. The new piece SAM 2 adds is memory: run a shared image encoder over every frame's pixels, and on frame 0 also pool the encoder's features over the prompted mask into a single memory embedding — a compact summary of "what the target looks like." For every later frame, compare that one memory embedding against every spatial location in the current frame's feature map (a dot-product correlation, producing a per-pixel similarity score) and threshold the result into a mask. Because this search is global — it scans the whole frame for the best match to the remembered appearance — it doesn't care where the object was in the previous frame, only what it looked like when it was last seen clearly.
Real SAM 2's image encoder is Hiera, a hierarchical Vision Transformer, so this lesson's encoder is a small Transformer too (Lesson 47's patchify-and-attend pattern) rather than a plain CNN. One real-architecture detail matters here more than it did for SAM's single-image case: Hiera restricts most of its attention to local windows rather than attending across the whole image at every layer. That's not just an efficiency trick — for this specific task, a patch's feature needs to describe its own local appearance so it can be matched against the same object in a later frame, and unrestricted global attention would blend in frame-varying context (like where the occluder happens to be) that has nothing to do with the target's identity. This lesson's encoder mirrors that with a small local attention window for the same reason.
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
def windowed_attn_mask(grid, window):
"""True where patch i may NOT attend to patch j -- outside a (window x window)
neighborhood in the 2D patch grid, mirroring Hiera's local-window attention. A
patch's feature should describe its own local appearance here, not get diluted by
unrelated, frame-varying global context (e.g. where the occluder happens to be)."""
coords = torch.stack(torch.meshgrid(torch.arange(grid), torch.arange(grid), indexing='ij'), dim=-1).reshape(-1, 2)
d = (coords[:, None, :] - coords[None, :, :]).abs().max(dim=-1).values
return d > (window // 2)
class Encoder(nn.Module):
def __init__(self, img_size=SIZE, patch_size=2, embed_dim=24, n_heads=4, n_layers=2, window=3):
super().__init__()
self.patch_size = patch_size
self.grid = img_size // patch_size
patch_dim = patch_size * patch_size
self.embed = nn.Linear(patch_dim, embed_dim)
self.register_buffer('pos_enc', positional_encoding(self.grid * self.grid, embed_dim))
self.register_buffer('attn_mask', windowed_attn_mask(self.grid, window))
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)
def forward(self, x):
p = patchify(x, self.patch_size)
tok = self.embed(p) + self.pos_enc
out = self.encoder(tok, mask=self.attn_mask)
B = x.shape[0]
feat = out.transpose(1, 2).reshape(B, -1, self.grid, self.grid)
return F.interpolate(feat, size=x.shape[-2:], mode='nearest') # (T, C, H, W), one feature map per frame
def correlate(feat, query):
# feat: (T, C, H, W); query: (T, C), broadcast per-frame -> (T, H, W) similarity map
T, C = query.shape
return (feat * query.view(T, C, 1, 1)).sum(dim=1)
def memory_embedding(feat0, mask0):
# pool frame 0's features over the prompted mask -> a single (C,) appearance vector
return (feat0 * mask0.unsqueeze(0)).sum(dim=(1, 2)) / (mask0.sum() + 1e-6)
N_CLIPS = 300
train_clips = [make_clip(np.random.default_rng(1000 + i)) for i in range(N_CLIPS)]
torch.manual_seed(0)
encoder = Encoder()
opt = torch.optim.Adam(encoder.parameters(), lr=0.003)
for epoch in range(1500):
frames, masks, occ, traj = train_clips[np.random.default_rng(epoch).integers(N_CLIPS)]
frames_t, masks_t = torch.tensor(frames).unsqueeze(1), torch.tensor(masks)
feat = encoder(frames_t)
query = memory_embedding(feat[0], masks_t[0])
corr = correlate(feat, query.unsqueeze(0).expand(frames_t.shape[0], -1))
loss = F.binary_cross_entropy_with_logits(corr, masks_t)
opt.zero_grad()
loss.backward()
opt.step()
print(f'final training loss: {loss.item():.4f}')
Does memory actually survive occlusion?¶
The real test is recovery: after the target disappears behind the occluder for a few frames, does the model find it again the instant it re-emerges? Compare the memory-based approach against a local tracker baseline in the spirit of Lesson 21's optical flow — it has no memory beyond the previous frame, only a small search window around the mask's last known position (a classic template-matching/KLT-style tracker). Both methods see the same frame-0 prompt; only how they use later frames differs.
def iou(a, b):
a, b = a.astype(bool), b.astype(bool)
inter, union = (a & b).sum(), (a | b).sum()
return inter / union if union > 0 else float('nan')
def eval_memory(clip):
frames, masks, occ, traj = clip
frames_t, masks_t = torch.tensor(frames).unsqueeze(1), torch.tensor(masks)
with torch.no_grad():
feat = encoder(frames_t)
query = memory_embedding(feat[0], masks_t[0])
corr = correlate(feat, query.unsqueeze(0).expand(frames_t.shape[0], -1))
preds = (torch.sigmoid(corr) > 0.5).numpy()
return [iou(preds[t], masks[t]) for t in range(len(frames))], preds
def eval_local_tracker(clip, search_radius=3):
frames, masks, occ, traj = clip
cy0, cx0 = np.array(np.where(masks[0] > 0)).mean(axis=1)
half = R + 1
template = frames[0][int(cy0) - half:int(cy0) + half, int(cx0) - half:int(cx0) + half].copy()
pos = np.array([cx0, cy0])
preds = [masks[0]]
for t in range(1, len(frames)):
img = frames[t]
best_score, best_pos = -1e9, pos.copy()
for dy in range(-search_radius, search_radius + 1):
for dx in range(-search_radius, search_radius + 1):
py, px = int(pos[1] + dy), int(pos[0] + dx)
if py - half < 0 or px - half < 0 or py + half > SIZE or px + half > SIZE:
continue
candidate = img[py - half:py + half, px - half:px + half]
score = -np.sum((candidate - template) ** 2) # best local match to the ORIGINAL template
if score > best_score:
best_score, best_pos = score, np.array([px, py])
pos = best_pos
yy, xx = np.mgrid[0:SIZE, 0:SIZE]
preds.append((((xx - pos[0]) ** 2 + (yy - pos[1]) ** 2) <= R ** 2).astype(np.float32))
return [iou(preds[t], masks[t]) for t in range(len(frames))], preds
test_clips = [make_clip(np.random.default_rng(5000 + i)) for i in range(80)]
mem_results = [eval_memory(c) for c in test_clips]
loc_results = [eval_local_tracker(c) for c in test_clips]
mem_ious = np.array([r[0] for r in mem_results])
loc_ious = np.array([r[0] for r in loc_results])
print('per-frame mean IoU, memory-based model: ', np.round(mem_ious.mean(axis=0), 3))
print('per-frame mean IoU, local-tracker baseline:', np.round(loc_ious.mean(axis=0), 3))
post1_mem = np.mean([mem_ious[i, c[2][-1] + 1] for i, c in enumerate(test_clips)])
post1_loc = np.mean([loc_ious[i, c[2][-1] + 1] for i, c in enumerate(test_clips)])
post2_mem = np.mean([mem_ious[i, c[2][-1] + 2] for i, c in enumerate(test_clips)])
post2_loc = np.mean([loc_ious[i, c[2][-1] + 2] for i, c in enumerate(test_clips)])
print(f'IoU 1 frame after occlusion ends: memory={post1_mem:.3f} local-tracker={post1_loc:.3f}')
print(f'IoU 2 frames after occlusion ends: memory={post2_mem:.3f} local-tracker={post2_loc:.3f}')
print(f'overall mean IoU (all frames): memory={mem_ious.mean():.3f} local-tracker={loc_ious.mean():.3f}')
fig, ax = plt.subplots(figsize=(7, 3.5))
ax.plot(mem_ious.mean(axis=0), 'o-', label='memory-based (global re-detection)')
ax.plot(loc_ious.mean(axis=0), 's-', label='local tracker (last-position search)')
ax.axvspan(2.5, 4.5, color='gray', alpha=0.2, label='typical occlusion window')
ax.set_xlabel('frame'); ax.set_ylabel('mean IoU vs. ground truth'); ax.legend(fontsize=8)
ax.set_title('Recovery after occlusion: memory vs. local tracking')
plt.show()
Right up to the occluder, both methods track the target well — there's no occlusion yet, so a local search window around the previous frame's position works just as well as a global memory search. The difference appears exactly where it should: the local tracker's last confident match before occlusion is near the occluder's edge, so once the target is hidden it has nothing good to latch onto, and its small search window can't find the target again once it re-emerges outside that window — it stays stuck near the occluder, and its IoU never fully recovers for the rest of the clip. The memory-based model has no such radius limit: it re-scans the entire frame for the embedding it stored back on frame 0, so the instant the target becomes visible again — anywhere in the frame — it's found. This is the core reason SAM 2-style memory beats classical flow-based tracking for real video segmentation: flow and local trackers accumulate drift and fail catastrophically under occlusion, while a memory of appearance (not position) survives it.
One number above is worth addressing directly rather than skating past: the overall mean IoU across every frame is not necessarily higher for the memory-based model, and can even come out lower. That's not a bug, and it isn't actually the number that matters here. During the occlusion itself, the target is genuinely invisible in the pixels — neither method can correctly segment something that isn't there, and the local tracker's "stay near the last known position" heuristic happens to overlap the ground truth a little by sheer geometric luck during those few frames (the target hasn't moved far from its last visible position yet), which inflates its average over the whole clip without it ever actually re-finding anything. A single clip-wide average blends "doing fine before occlusion," "both methods failing during it," and "one method recovering after it" into one number, so it can hide the exact effect this lesson is about. The two numbers that actually isolate that effect — IoU one and two frames after the target reappears — are the ones to trust, and they show the real story cleanly: this is a reminder that choosing the right evaluation metric for what you actually want to measure matters as much as the model itself.
show_frames = [0, 3, 4, 5, 6, 8]
fig, axes = plt.subplots(3, len(show_frames), figsize=(10, 5))
ex_frames, ex_masks, ex_occ, ex_traj = test_clips[1]
_, ex_mem_preds = eval_memory(test_clips[1])
_, ex_loc_preds = eval_local_tracker(test_clips[1])
for col, t in enumerate(show_frames):
axes[0, col].imshow(ex_frames[t], cmap='gray', vmin=0, vmax=1)
axes[0, col].set_title(f't={t}' + (' (occ)' if t in ex_occ else ''), fontsize=9)
axes[1, col].imshow(ex_mem_preds[t], cmap='gray', vmin=0, vmax=1)
axes[2, col].imshow(ex_loc_preds[t], cmap='gray', vmin=0, vmax=1)
for row in range(3):
axes[row, col].axis('off')
axes[0, 0].set_title(f't=0', fontsize=9, loc='left')
fig.text(0.02, 0.78, 'frame', fontsize=9, ha='left')
fig.text(0.02, 0.5, 'memory\nprediction', fontsize=9, ha='left')
fig.text(0.02, 0.22, 'local-tracker\nprediction', fontsize=9, ha='left')
plt.suptitle('One example clip: memory recovers immediately, the local tracker stays lost', y=1.02)
plt.tight_layout()
plt.show()
From a toy memory bank to SAM 2¶
This lesson's memory is a single, fixed embedding frozen at the moment of the prompt. SAM 2 (Ravi et al., 2024) extends the same idea with a genuine memory bank: a short rolling window of recent frames' embeddings (not just frame 0), so the remembered appearance can drift gracefully as the object's pose, lighting, or viewpoint changes over a long video, and a streaming memory attention module that lets each new frame attend over that whole bank rather than a single vector. The core mechanism this lesson demonstrates — re-finding an object by appearance rather than by position — is exactly what makes that scale to real, minutes-long video reliably.
Exercises¶
- Shrink
search_radiusineval_local_trackerfrom 3 to 1, and separately grow it to 8. How does the local tracker's post-occlusion IoU change in each direction, and is there a radius large enough to match the memory-based model's recovery — what would that cost in computation per frame? - This lesson's memory is a single embedding pooled once, on frame 0, and never updated. Modify
memory_embedding(or the training loop) to instead average the pooled embedding over the last 3 confidently-visible frames (updating the running average each time a frame's prediction is confident) rather than only frame 0. Does this make recovery more or less robust if the object's appearance changes partway through the clip (e.g., try varyingOCC_HALFor adding per-frame intensity noise)? - The occluder in this lesson is static and always in the same place. Change
make_clipso the occluder itself also moves across frames. Does the memory-based model still recover cleanly, or does a moving occluder expose a weakness in a single-frame, position-blind correlation search (e.g., could the occluder itself ever get mistaken for the target)?