Lesson 57: Open-Vocabulary Segmentation¶
Lesson 53's Segment Anything segments on demand from a spatial prompt (a point or box) but has no notion of what it's segmenting — no text, no categories. Lesson 56's open-vocabulary detection goes the other way: it localizes objects from arbitrary text categories, but only as boxes. Open-vocabulary segmentation combines both threads: prompt the model with a text concept ("a player in white", "the cracked tile") and get back pixel-accurate masks for every matching instance, in the style of SAM 3 (Carion et al., 2025) — a text-and-image encoder pair (CLIP-style, Lesson 55) grounding the concept, feeding a detector/segmenter that produces instance masks with identities, not just a single yes/no match.
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
Scenes with shape and color¶
Scenes now carry two independent attributes per instance — a shape (plus or circle) and a color (red or blue) — so a text concept has to combine both to pick out the right instances. This is what makes "open-vocabulary" meaningfully harder than Lesson 56's single-attribute category names: the model has to genuinely compose two words into one visual criterion, not just match a single learned class name.
SIZE = 32
SHAPES = ['plus', 'circle']
COLORS = ['red', 'blue']
COLOR_RGB = {'red': np.array([1.0, 0.15, 0.15]), 'blue': np.array([0.15, 0.3, 1.0])}
def draw_shape(canvas, shape_type, color, cx, cy, r=5):
size = canvas.shape[0]
if shape_type == 'plus':
m = np.zeros((size, size), dtype=bool)
m[max(cy-1, 0):cy+2, max(cx-r, 0):cx+r+1] = True
m[max(cy-r, 0):cy+r+1, max(cx-1, 0):cx+2] = True
else:
yy, xx = np.mgrid[0:size, 0:size]
m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2
canvas[m] = COLOR_RGB[color]
return m
def make_scene(rng, size=SIZE, n_shapes=4, r=5):
"""A scene of n_shapes instances, each a random (shape, color) combo, placed to avoid heavy overlap."""
canvas = np.zeros((size, size, 3), dtype=np.float32)
instances = [] # (shape, color, mask, cx, cy)
attempts = 0
while len(instances) < n_shapes and attempts < 40:
attempts += 1
shape, color = rng.choice(SHAPES), rng.choice(COLORS)
cx, cy = rng.integers(r + 3, size - r - 3), rng.integers(r + 3, size - r - 3)
if any((cx - icx) ** 2 + (cy - icy) ** 2 < (2 * r) ** 2 for (_, _, _, icx, icy) in instances):
continue
m = draw_shape(canvas, shape, color, cx, cy, r)
instances.append((shape, color, m, cx, cy))
canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)
return canvas, instances
rng = np.random.default_rng(7)
fig, axes = plt.subplots(1, 4, figsize=(9, 2.5))
for ax in axes:
canvas, instances = make_scene(rng)
ax.imshow(canvas)
ax.axis('off')
plt.suptitle('Scenes: each instance is an independent (shape, color) pair')
plt.show()
Grounding a text concept in the image, pixel by pixel¶
Lesson 55's CLIP pools a whole image down to one vector and compares it against one text vector — fine for "does this image match this caption," useless for "which pixels match this caption." Open-vocabulary segmentation needs the comparison to happen at every pixel independently, so the image encoder here is a ViT (Lesson 47's patchify + Transformer-encoder pattern) that keeps its per-patch tokens instead of pooling them into one vector: each patch becomes a coarse (D, grid, grid) feature map, which is then upsampled back to full pixel resolution — exactly how real ViT-based dense-prediction heads (segmentation, depth) work, since attention naturally runs on a sequence of patches, not individual pixels. The text encoder is now a small self-attention Transformer too, same pattern as Lesson 56's, pooled to one vector per concept.
Grounding a concept is then just a per-pixel cosine similarity between that one text vector and every pixel's feature vector: broadcast the text embedding against the whole feature map and take a dot product at each spatial location. This is exactly the fusion step real open-vocabulary segmenters use to connect language and pixels, and it is worth being precise about what it is not: there is no generation happening here, nothing is produced one step at a time, and nothing attends to its own previous output. It's a similarity lookup, computed in a single forward pass — the same encoder-only spirit as CLIP, just evaluated everywhere in the image instead of once.
VOCAB = ['a', 'photo', 'of', 'red', 'blue', 'plus', 'circle']
W2ID = {w: i for i, w in enumerate(VOCAB)}
def caption_for(color, shape):
return [W2ID[w] for w in ['a', 'photo', 'of', color, shape]]
def patchify(img, patch_size):
B, C, H, W = img.shape
P = patch_size
patches = img.unfold(2, P, P).unfold(3, P, P) # (B, C, H/P, W/P, 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) # (B, num_patches, 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 ImageEncoder(nn.Module):
"""ViT: patchify into a coarse grid, run a Transformer encoder over the patch tokens,
then upsample the per-patch features back to full pixel resolution -- dense per-pixel
output, not a pooled single vector."""
def __init__(self, img_size=SIZE, patch_size=4, out_dim=16, embed_dim=32, n_heads=4, n_layers=2):
super().__init__()
self.patch_size = patch_size
self.grid = img_size // patch_size
patch_dim = 3 * 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))
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.proj = nn.Linear(embed_dim, out_dim)
def forward(self, x):
B = x.shape[0]
p = patchify(x, self.patch_size) # (B, grid*grid, 3*P*P)
tok = self.embed(p) + self.pos_enc # (B, grid*grid, embed_dim)
out = self.proj(self.encoder(tok)) # (B, grid*grid, out_dim)
feat = out.transpose(1, 2).reshape(B, -1, self.grid, self.grid) # (B, out_dim, grid, grid)
return F.interpolate(feat, size=x.shape[-2:], mode='bilinear', align_corners=False)
class TextEncoder(nn.Module):
"""A small self-attention text encoder over the caption tokens, mean-pooled to one
vector per concept -- same pattern as Lesson 56's."""
def __init__(self, vocab_size, out_dim=16, embed_dim=16, n_heads=2, n_layers=1, max_len=8):
super().__init__()
self.tok_embed = nn.Embedding(vocab_size, embed_dim)
self.pos_embed = nn.Embedding(max_len, 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.proj = nn.Linear(embed_dim, out_dim)
def forward(self, ids):
B, T = ids.shape
pos = torch.arange(T, device=ids.device)
x = self.tok_embed(ids) + self.pos_embed(pos)[None]
out = self.encoder(x)
return self.proj(out.mean(dim=1))
def ground(img_feat, txt_emb):
"""Per-pixel cosine similarity between a feature map and one text concept embedding."""
img_n = F.normalize(img_feat, dim=1)
txt_n = F.normalize(txt_emb, dim=1)
return torch.einsum('bdhw,bd->bhw', img_n, txt_n)
def build_dataset(rng, n_scenes):
"""One training example per (scene, concept actually present in that scene) --
everything else in the scene becomes a negative pixel for free, since only the
queried concept's instances are marked 1 in the target mask."""
scenes, examples = [], []
for si in range(n_scenes):
canvas, instances = make_scene(rng)
scenes.append(canvas)
present = set((color, shape) for (shape, color, m, cx, cy) in instances)
for (color, shape) in present:
target = np.zeros((SIZE, SIZE), dtype=np.float32)
for (s, c, m, cx, cy) in instances:
if s == shape and c == color:
target[m] = 1.0
examples.append((si, caption_for(color, shape), target))
return np.array(scenes, dtype=np.float32), examples
rng = np.random.default_rng(7)
scenes_tr, examples_tr = build_dataset(rng, 500)
canvases_t = torch.tensor(scenes_tr).permute(0, 3, 1, 2) # (N, 3, H, W)
print(f'{len(scenes_tr)} scenes, {len(examples_tr)} (scene, concept) training examples')
torch.manual_seed(0)
img_enc = ImageEncoder()
txt_enc = TextEncoder(len(VOCAB))
opt = torch.optim.Adam(list(img_enc.parameters()) + list(txt_enc.parameters()), lr=0.005)
for epoch in range(800):
idx = np.random.default_rng(epoch).permutation(len(examples_tr))[:64]
batch_si = torch.tensor([examples_tr[i][0] for i in idx])
batch_caps = torch.tensor([examples_tr[i][1] for i in idx])
batch_tgt = torch.tensor(np.stack([examples_tr[i][2] for i in idx]))
feat = img_enc(canvases_t[batch_si])
txt = txt_enc(batch_caps)
scores = ground(feat, txt)
loss = F.binary_cross_entropy_with_logits(scores * 5.0, batch_tgt)
opt.zero_grad()
loss.backward()
opt.step()
print(f'final training loss: {loss.item():.4f}')
From a concept's score map to per-instance masks¶
Thresholding the grounding score gives one combined mask for every instance matching the concept at once — if a scene has two red circles, both light up together, with no notion that they're two separate objects with two separate identities. SAM 3 (Carion et al., 2025) solves this with a DETR-style detector: a transformer decoder (in the generic sense — cross-attention combining a fixed set of learned "object slot" queries with the grounded image features) that outputs a whole set of instance masks and identities in one forward pass. It is worth being just as precise here as above about what kind of decoder this is, because this course has one other decoder nearby that works completely differently: SAM 3's detector decoder is non-autoregressive — it predicts its whole set of output masks simultaneously, with no step where it looks at its own previous output. That is fundamentally different from the autoregressive decoder Lesson 59 builds, which generates one word at a time, each new word conditioned on every word generated so far. Both are legitimately called "a Transformer decoder" in the literature, and conflating the two is one of the easiest mistakes to make with this vocabulary.
This lesson's toy version skips the learned decoder entirely and separates instances the classical way instead: threshold the score map, then run connected-component labeling (Lesson 4's flood fill, via cv2.connectedComponents) on the result. It recovers the same outcome — one label per instance — without a trained decoder at all, which is a fair toy stand-in only because this lesson's shapes never touch; SAM 3's learned decoder is doing real work precisely in the cases (touching or overlapping instances) where connected components would fail.
def make_two_same_concept_scene(rng, size=SIZE, r=5):
"""Two red circles plus one blue-plus distractor, so one concept has two true instances."""
canvas = np.zeros((size, size, 3), dtype=np.float32)
positions = [(8, 8), (24, 24), (24, 8)]
rng.shuffle(positions)
combos = [('circle', 'red'), ('circle', 'red'), ('plus', 'blue')]
instances = []
for (shape, color), (cx, cy) in zip(combos, positions):
m = draw_shape(canvas, shape, color, cx, cy, r)
instances.append((shape, color, m, cx, cy))
canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)
return canvas, instances
demo_rng = np.random.default_rng(321)
canvas, instances = make_two_same_concept_scene(demo_rng)
true_count = sum(1 for (s, c, m, cx, cy) in instances if s == 'circle' and c == 'red')
with torch.no_grad():
feat = img_enc(torch.tensor(canvas).permute(2, 0, 1).unsqueeze(0))
txt = txt_enc(torch.tensor([caption_for('red', 'circle')]))
score = ground(feat, txt)[0].numpy()
concept_mask = (score > 0.0).astype(np.uint8)
n_components, labels = cv2.connectedComponents(concept_mask)
print(f'"red circle" query: {n_components - 1} separate instance(s) found (true count: {true_count})')
fig, axes = plt.subplots(1, 3, figsize=(8, 2.8))
axes[0].imshow(canvas); axes[0].set_title('scene'); axes[0].axis('off')
axes[1].imshow(concept_mask, cmap='gray'); axes[1].set_title('"red circle" score, thresholded'); axes[1].axis('off')
axes[2].imshow(labels, cmap='tab10'); axes[2].set_title('connected components\n(one color per instance)'); axes[2].axis('off')
plt.show()
Disambiguating closely related concepts¶
The real test of a compositional concept, rather than a single learned class name: does the model actually use both words? Build a scene containing one instance of all four (shape, color) combinations, query it with the correct concept ("red circle"), then query the exact same scene with a color near-miss ("blue circle" — right shape, wrong color) and a shape near-miss ("red plus" — right color, wrong shape), and check all three predictions against the one true "red circle" region.
def make_disambig_scene(rng, size=SIZE, r=5):
"""All four (shape, color) combinations present exactly once."""
canvas = np.zeros((size, size, 3), dtype=np.float32)
positions = [(8, 8), (24, 8), (8, 24), (24, 24)]
rng.shuffle(positions)
combos = [('circle', 'red'), ('circle', 'blue'), ('plus', 'red'), ('plus', 'blue')]
instances = []
for (shape, color), (cx, cy) in zip(combos, positions):
m = draw_shape(canvas, shape, color, cx, cy, r)
instances.append((shape, color, m, cx, cy))
canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)
return canvas, instances
def precision_recall(pred_mask, true_mask):
pred, true = pred_mask.astype(bool), true_mask.astype(bool)
tp = (pred & true).sum()
prec = tp / pred.sum() if pred.sum() > 0 else float('nan')
rec = tp / true.sum() if true.sum() > 0 else float('nan')
return prec, rec
queries = {
'correct concept "red circle"': ('red', 'circle'),
'color near-miss "blue circle"': ('blue', 'circle'),
'shape near-miss "red plus"': ('red', 'plus'),
}
results = {name: {'prec': [], 'rec': []} for name in queries}
eval_rng = np.random.default_rng(99)
with torch.no_grad():
for _ in range(80):
canvas, instances = make_disambig_scene(eval_rng)
feat = img_enc(torch.tensor(canvas).permute(2, 0, 1).unsqueeze(0))
true_mask = next(m for (s, c, m, cx, cy) in instances if s == 'circle' and c == 'red')
for name, (color, shape) in queries.items():
txt = txt_enc(torch.tensor([caption_for(color, shape)]))
pred_mask = ground(feat, txt)[0].numpy() > 0.0
prec, rec = precision_recall(pred_mask, true_mask)
results[name]['prec'].append(prec)
results[name]['rec'].append(rec)
print('Each query evaluated against the scene\'s one true "red circle" region:')
for name, r in results.items():
print(f' {name:<32} precision={np.nanmean(r["prec"]):.3f} recall={np.nanmean(r["rec"]):.3f}')
Precision and recall both collapse to essentially zero for both near-miss queries, while the correct concept recovers the true region almost perfectly — the model isn't matching on shape alone or color alone, it's grounding the full composed phrase, and genuinely suppresses an instance that satisfies only half of it. That compositional behavior, not just "detect a fixed list of categories," is the actual claim behind "open-vocabulary": the model never saw the phrase "red circle" paired with this image during training, only the general pattern that a concept's words jointly determine which pixels light up.
This closes the arc this subsection started: Lesson 53 made "where" an explicit input instead of something to infer, Lesson 56 made "what" an open text vocabulary instead of a fixed label set, and this lesson combines both into one query. What's still missing — and what Lesson 59 picks up next — is anything generative: every model in this lesson answers a question that's already fully specified by its input (which pixels match this fixed phrase), never composes a novel sentence of its own.
Exercises¶
- Add a third color (e.g.
'green') toCOLORSand retrain. Does 3-color disambiguation hold up as cleanly as the 2-color case, or does precision/recall degrade on the near-miss queries? - The grounding score in
ground()is an unweighted cosine similarity between the full 16-dim text and image vectors. Try training with a shorter embedding dimension (e.g.out_dim=4in both encoders) — does disambiguation between color and shape still work, or does a too-small embedding space start blending the two attributes together? make_two_same_concept_scenerelies on the two red circles never touching, so connected-component labeling can tell them apart. Modify it to place the two red circles close enough to touch (e.g. 6 pixels apart center-to-center instead of using the fixedpositionslist) and rerun the instance-counting demo. Doescv2.connectedComponentsstill report 2 instances, and does this failure mode make it clearer why SAM 3 uses a learned decoder here instead of connected components?