{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Lesson 57: Open-Vocabulary Segmentation\n\nLesson {{lesson:segment_anything}}'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 {{lesson:open_vocabulary_detection}}'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 {{lesson:vision_language_models}}) grounding the concept, feeding a detector/segmenter that produces instance masks with identities, not just a single yes/no match."
},
{
"cell_type": "code",
"id": "6dc554ba",
"source": "import numpy as np\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "03e4ec13",
"source": "## Scenes with shape *and* color\n\nScenes 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 {{lesson:open_vocabulary_detection}}'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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "23f74587",
"source": "SIZE = 32\nSHAPES = ['plus', 'circle']\nCOLORS = ['red', 'blue']\nCOLOR_RGB = {'red': np.array([1.0, 0.15, 0.15]), 'blue': np.array([0.15, 0.3, 1.0])}\n\n\ndef draw_shape(canvas, shape_type, color, cx, cy, r=5):\n size = canvas.shape[0]\n if shape_type == 'plus':\n m = np.zeros((size, size), dtype=bool)\n m[max(cy-1, 0):cy+2, max(cx-r, 0):cx+r+1] = True\n m[max(cy-r, 0):cy+r+1, max(cx-1, 0):cx+2] = True\n else:\n yy, xx = np.mgrid[0:size, 0:size]\n m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2\n canvas[m] = COLOR_RGB[color]\n return m\n\n\ndef make_scene(rng, size=SIZE, n_shapes=4, r=5):\n \"\"\"A scene of n_shapes instances, each a random (shape, color) combo, placed to avoid heavy overlap.\"\"\"\n canvas = np.zeros((size, size, 3), dtype=np.float32)\n instances = [] # (shape, color, mask, cx, cy)\n attempts = 0\n while len(instances) < n_shapes and attempts < 40:\n attempts += 1\n shape, color = rng.choice(SHAPES), rng.choice(COLORS)\n cx, cy = rng.integers(r + 3, size - r - 3), rng.integers(r + 3, size - r - 3)\n if any((cx - icx) ** 2 + (cy - icy) ** 2 < (2 * r) ** 2 for (_, _, _, icx, icy) in instances):\n continue\n m = draw_shape(canvas, shape, color, cx, cy, r)\n instances.append((shape, color, m, cx, cy))\n canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)\n return canvas, instances\n\n\nrng = np.random.default_rng(7)\nfig, axes = plt.subplots(1, 4, figsize=(9, 2.5))\nfor ax in axes:\n canvas, instances = make_scene(rng)\n ax.imshow(canvas)\n ax.axis('off')\nplt.suptitle('Scenes: each instance is an independent (shape, color) pair')\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "b9cc08b4",
"source": "## Grounding a text concept in the image, pixel by pixel\n\nLesson {{lesson:vision_language_models}}'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 {{lesson:vision_transformers}}'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 {{lesson:open_vocabulary_detection}}'s, pooled to one vector per concept.\n\nGrounding 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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "b2d6523e",
"source": "VOCAB = ['a', 'photo', 'of', 'red', 'blue', 'plus', 'circle']\nW2ID = {w: i for i, w in enumerate(VOCAB)}\n\ndef caption_for(color, shape):\n return [W2ID[w] for w in ['a', 'photo', 'of', color, shape]]\n\n\ndef patchify(img, patch_size):\n B, C, H, W = img.shape\n P = patch_size\n patches = img.unfold(2, P, P).unfold(3, P, P) # (B, C, H/P, W/P, P, P)\n patches = patches.contiguous().view(B, C, -1, P, P).permute(0, 2, 1, 3, 4)\n return patches.reshape(B, -1, C * P * P) # (B, num_patches, C*P*P)\n\n\ndef positional_encoding(T, D):\n pos = torch.arange(T).unsqueeze(1).float()\n i = torch.arange(D).unsqueeze(0).float()\n angle_rates = 1.0 / (10000 ** (2 * (i // 2) / D))\n angles = pos * angle_rates\n pe = torch.zeros(T, D)\n pe[:, 0::2] = torch.sin(angles[:, 0::2])\n pe[:, 1::2] = torch.cos(angles[:, 1::2])\n return pe\n\n\nclass ImageEncoder(nn.Module):\n \"\"\"ViT: patchify into a coarse grid, run a Transformer encoder over the patch tokens,\n then upsample the per-patch features back to full pixel resolution -- dense per-pixel\n output, not a pooled single vector.\"\"\"\n def __init__(self, img_size=SIZE, patch_size=4, out_dim=16, embed_dim=32, n_heads=4, n_layers=2):\n super().__init__()\n self.patch_size = patch_size\n self.grid = img_size // patch_size\n patch_dim = 3 * patch_size * patch_size\n self.embed = nn.Linear(patch_dim, embed_dim)\n self.register_buffer('pos_enc', positional_encoding(self.grid * self.grid, embed_dim))\n layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n self.proj = nn.Linear(embed_dim, out_dim)\n\n def forward(self, x):\n B = x.shape[0]\n p = patchify(x, self.patch_size) # (B, grid*grid, 3*P*P)\n tok = self.embed(p) + self.pos_enc # (B, grid*grid, embed_dim)\n out = self.proj(self.encoder(tok)) # (B, grid*grid, out_dim)\n feat = out.transpose(1, 2).reshape(B, -1, self.grid, self.grid) # (B, out_dim, grid, grid)\n return F.interpolate(feat, size=x.shape[-2:], mode='bilinear', align_corners=False)\n\n\nclass TextEncoder(nn.Module):\n \"\"\"A small self-attention text encoder over the caption tokens, mean-pooled to one\n vector per concept -- same pattern as Lesson {{lesson:open_vocabulary_detection}}'s.\"\"\"\n def __init__(self, vocab_size, out_dim=16, embed_dim=16, n_heads=2, n_layers=1, max_len=8):\n super().__init__()\n self.tok_embed = nn.Embedding(vocab_size, embed_dim)\n self.pos_embed = nn.Embedding(max_len, embed_dim)\n layer = nn.TransformerEncoderLayer(embed_dim, n_heads, dim_feedforward=embed_dim * 2,\n batch_first=True, dropout=0.0)\n self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)\n self.proj = nn.Linear(embed_dim, out_dim)\n\n def forward(self, ids):\n B, T = ids.shape\n pos = torch.arange(T, device=ids.device)\n x = self.tok_embed(ids) + self.pos_embed(pos)[None]\n out = self.encoder(x)\n return self.proj(out.mean(dim=1))\n\n\ndef ground(img_feat, txt_emb):\n \"\"\"Per-pixel cosine similarity between a feature map and one text concept embedding.\"\"\"\n img_n = F.normalize(img_feat, dim=1)\n txt_n = F.normalize(txt_emb, dim=1)\n return torch.einsum('bdhw,bd->bhw', img_n, txt_n)\n\n\ndef build_dataset(rng, n_scenes):\n \"\"\"One training example per (scene, concept actually present in that scene) --\n everything else in the scene becomes a negative pixel for free, since only the\n queried concept's instances are marked 1 in the target mask.\"\"\"\n scenes, examples = [], []\n for si in range(n_scenes):\n canvas, instances = make_scene(rng)\n scenes.append(canvas)\n present = set((color, shape) for (shape, color, m, cx, cy) in instances)\n for (color, shape) in present:\n target = np.zeros((SIZE, SIZE), dtype=np.float32)\n for (s, c, m, cx, cy) in instances:\n if s == shape and c == color:\n target[m] = 1.0\n examples.append((si, caption_for(color, shape), target))\n return np.array(scenes, dtype=np.float32), examples\n\n\nrng = np.random.default_rng(7)\nscenes_tr, examples_tr = build_dataset(rng, 500)\ncanvases_t = torch.tensor(scenes_tr).permute(0, 3, 1, 2) # (N, 3, H, W)\nprint(f'{len(scenes_tr)} scenes, {len(examples_tr)} (scene, concept) training examples')\n\ntorch.manual_seed(0)\nimg_enc = ImageEncoder()\ntxt_enc = TextEncoder(len(VOCAB))\nopt = torch.optim.Adam(list(img_enc.parameters()) + list(txt_enc.parameters()), lr=0.005)\n\nfor epoch in range(800):\n idx = np.random.default_rng(epoch).permutation(len(examples_tr))[:64]\n batch_si = torch.tensor([examples_tr[i][0] for i in idx])\n batch_caps = torch.tensor([examples_tr[i][1] for i in idx])\n batch_tgt = torch.tensor(np.stack([examples_tr[i][2] for i in idx]))\n\n feat = img_enc(canvases_t[batch_si])\n txt = txt_enc(batch_caps)\n scores = ground(feat, txt)\n loss = F.binary_cross_entropy_with_logits(scores * 5.0, batch_tgt)\n opt.zero_grad()\n loss.backward()\n opt.step()\n\nprint(f'final training loss: {loss.item():.4f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "4cc4f006",
"source": "## From a concept's score map to per-instance masks\n\nThresholding 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 {{lesson:visual_question_answering}} 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.\n\nThis 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 {{lesson:floodfill_connected_components}}'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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "f2552c85",
"source": "def make_two_same_concept_scene(rng, size=SIZE, r=5):\n \"\"\"Two red circles plus one blue-plus distractor, so one concept has two true instances.\"\"\"\n canvas = np.zeros((size, size, 3), dtype=np.float32)\n positions = [(8, 8), (24, 24), (24, 8)]\n rng.shuffle(positions)\n combos = [('circle', 'red'), ('circle', 'red'), ('plus', 'blue')]\n instances = []\n for (shape, color), (cx, cy) in zip(combos, positions):\n m = draw_shape(canvas, shape, color, cx, cy, r)\n instances.append((shape, color, m, cx, cy))\n canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)\n return canvas, instances\n\ndemo_rng = np.random.default_rng(321)\ncanvas, instances = make_two_same_concept_scene(demo_rng)\ntrue_count = sum(1 for (s, c, m, cx, cy) in instances if s == 'circle' and c == 'red')\n\nwith torch.no_grad():\n feat = img_enc(torch.tensor(canvas).permute(2, 0, 1).unsqueeze(0))\n txt = txt_enc(torch.tensor([caption_for('red', 'circle')]))\n score = ground(feat, txt)[0].numpy()\n\nconcept_mask = (score > 0.0).astype(np.uint8)\nn_components, labels = cv2.connectedComponents(concept_mask)\nprint(f'\"red circle\" query: {n_components - 1} separate instance(s) found (true count: {true_count})')\n\nfig, axes = plt.subplots(1, 3, figsize=(8, 2.8))\naxes[0].imshow(canvas); axes[0].set_title('scene'); axes[0].axis('off')\naxes[1].imshow(concept_mask, cmap='gray'); axes[1].set_title('\"red circle\" score, thresholded'); axes[1].axis('off')\naxes[2].imshow(labels, cmap='tab10'); axes[2].set_title('connected components\\n(one color per instance)'); axes[2].axis('off')\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "9e04b205",
"source": "## Disambiguating closely related concepts\n\nThe 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.",
"metadata": {}
},
{
"cell_type": "code",
"id": "ea36dc4f",
"source": "def make_disambig_scene(rng, size=SIZE, r=5):\n \"\"\"All four (shape, color) combinations present exactly once.\"\"\"\n canvas = np.zeros((size, size, 3), dtype=np.float32)\n positions = [(8, 8), (24, 8), (8, 24), (24, 24)]\n rng.shuffle(positions)\n combos = [('circle', 'red'), ('circle', 'blue'), ('plus', 'red'), ('plus', 'blue')]\n instances = []\n for (shape, color), (cx, cy) in zip(combos, positions):\n m = draw_shape(canvas, shape, color, cx, cy, r)\n instances.append((shape, color, m, cx, cy))\n canvas = np.clip(canvas + rng.normal(0, 0.03, canvas.shape), 0, 1).astype(np.float32)\n return canvas, instances\n\n\ndef precision_recall(pred_mask, true_mask):\n pred, true = pred_mask.astype(bool), true_mask.astype(bool)\n tp = (pred & true).sum()\n prec = tp / pred.sum() if pred.sum() > 0 else float('nan')\n rec = tp / true.sum() if true.sum() > 0 else float('nan')\n return prec, rec\n\n\nqueries = {\n 'correct concept \"red circle\"': ('red', 'circle'),\n 'color near-miss \"blue circle\"': ('blue', 'circle'),\n 'shape near-miss \"red plus\"': ('red', 'plus'),\n}\nresults = {name: {'prec': [], 'rec': []} for name in queries}\n\neval_rng = np.random.default_rng(99)\nwith torch.no_grad():\n for _ in range(80):\n canvas, instances = make_disambig_scene(eval_rng)\n feat = img_enc(torch.tensor(canvas).permute(2, 0, 1).unsqueeze(0))\n true_mask = next(m for (s, c, m, cx, cy) in instances if s == 'circle' and c == 'red')\n\n for name, (color, shape) in queries.items():\n txt = txt_enc(torch.tensor([caption_for(color, shape)]))\n pred_mask = ground(feat, txt)[0].numpy() > 0.0\n prec, rec = precision_recall(pred_mask, true_mask)\n results[name]['prec'].append(prec)\n results[name]['rec'].append(rec)\n\nprint('Each query evaluated against the scene\\'s one true \"red circle\" region:')\nfor name, r in results.items():\n print(f' {name:<32} precision={np.nanmean(r[\"prec\"]):.3f} recall={np.nanmean(r[\"rec\"]):.3f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "9739f183",
"source": "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.\n\nThis closes the arc this subsection started: Lesson {{lesson:segment_anything}} made \"where\" an explicit input instead of something to infer, Lesson {{lesson:open_vocabulary_detection}} 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 {{lesson:visual_question_answering}} 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.\n\n### Exercises\n\n1. Add a third color (e.g. `'green'`) to `COLORS` and retrain. Does 3-color disambiguation hold up as cleanly as the 2-color case, or does precision/recall degrade on the near-miss queries?\n2. 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=4` in both encoders) — does disambiguation between color and shape still work, or does a too-small embedding space start blending the two attributes together?\n3. `make_two_same_concept_scene` relies 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 fixed `positions` list) and rerun the instance-counting demo. Does `cv2.connectedComponents` still report 2 instances, and does this failure mode make it clearer why SAM 3 uses a learned decoder here instead of connected components?",
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.x"
}
},
"nbformat": 4,
"nbformat_minor": 5
}