{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9138458",
   "metadata": {},
   "source": "# Lesson 56: Open-Vocabulary Detection\n\nLesson 41's detector predicts one of a small, fixed set of classes baked in at training time — exactly Lesson {{lesson:evaluating_a_classifier}}'s classification limitation, just applied to boxes instead of whole images. **Open-vocabulary detection** (Grounding DINO, OWL-ViT) removes that limit by splitting detection into two separable skills: *finding* objects (a class-agnostic skill, surprisingly transferable across object categories) and *naming* them (delegated to a CLIP-style text-embedding match, Lesson {{lesson:vision_language_models}}, which can recognize any category describable in words). This lesson builds both halves and tests the combination on a category the detector's box regressor never received a single labeled box for."
  },
  {
   "cell_type": "code",
   "id": "97c114a2",
   "source": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "7b2d17f6",
   "source": "## Base classes vs. a novel class\n\nThree shapes: plus, circle, square. The box regressor will only ever see labeled boxes for two of them — the **base classes**. Square is the **novel class**: the detector never sees a single box-labeled square during training, only plus and circle.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "f8b8572b",
   "source": "SIZE = 32\nSHAPES = ['plus', 'circle', 'square']\nBASE_SHAPES = ['plus', 'circle']\nNOVEL_SHAPE = 'square'\n\ndef make_image(shape_type, cx, cy, size=SIZE, r=6):\n    img = np.zeros((size, size), dtype=np.float32)\n    if shape_type == 'plus':\n        img[cy-1:cy+2, cx-r:cx+r+1] = 1.0\n        img[cy-r:cy+r+1, cx-1:cx+2] = 1.0\n    elif shape_type == 'circle':\n        yy, xx = np.mgrid[0:size, 0:size]\n        img[((xx-cx)**2 + (yy-cy)**2) <= r**2] = 1.0\n    else:  # square\n        img[cy-r:cy+r+1, cx-r:cx+r+1] = 1.0\n    return img\n\ndef make_detection_example(rng, shapes, size=SIZE, r=6):\n    shape_type = rng.choice(shapes)\n    cx, cy = rng.integers(r + 2, size - r - 2), rng.integers(r + 2, size - r - 2)\n    img = make_image(shape_type, cx, cy)\n    img = np.clip(img + rng.normal(0, 0.05, img.shape), 0, 1).astype(np.float32)\n    box = np.array([cx - r, cy - r, 2 * r, 2 * r], dtype=np.float32) / size\n    return img, box, shape_type\n\ndef crop_and_resize(img, box, out_size=SIZE):\n    x0, y0, w, h = (box * SIZE).astype(int)\n    x0, y0 = max(0, x0), max(0, y0)\n    x1, y1 = min(SIZE, x0 + max(w, 1)), min(SIZE, y0 + max(h, 1))\n    crop = img[y0:y1, x0:x1]\n    if crop.size == 0:\n        return np.zeros((out_size, out_size), dtype=np.float32)\n    resized = F.interpolate(torch.tensor(crop[None, None]), size=(out_size, out_size),\n                             mode='bilinear', align_corners=False)\n    return resized[0, 0].numpy()\n\nfig, axes = plt.subplots(1, 3, figsize=(7, 2.5))\ndemo_rng = np.random.default_rng(42)\nfor ax, shape in zip(axes, SHAPES):\n    im, box, _ = make_detection_example(demo_rng, [shape])\n    ax.imshow(im, cmap='gray')\n    x0, y0, w, h = box * SIZE\n    ax.add_patch(patches.Rectangle((x0, y0), w, h, edgecolor='lime', facecolor='none', linewidth=2))\n    ax.set_title(f'{shape}  (base={shape in BASE_SHAPES})', fontsize=9)\n    ax.axis('off')\nplt.show()",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "c2b44700",
   "source": "## Half 1: a class-agnostic box regressor\n\nA ViT-style box regressor: Lesson {{lesson:vision_transformers}}'s patchify + positional-encoding + Transformer-encoder + `[CLS]`-token pattern, with the classification head swapped for a 4-number box head — real open-vocabulary detectors (OWL-ViT, Grounding DINO) localize with a ViT backbone too, not a plain CNN. Only the *target* changes from Lesson {{lesson:object_detection_regression}}'s box regression: never a class label, only \"where is the object,\" trained exclusively on plus and circle scenes.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "4d0507fe",
   "source": "def 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\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\nclass BoxRegressor(nn.Module):\n    def __init__(self, img_size=SIZE, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):\n        super().__init__()\n        self.patch_size = patch_size\n        n_patches = (img_size // patch_size) ** 2\n        patch_dim = 1 * patch_size * patch_size\n        self.embed = nn.Linear(patch_dim, embed_dim)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)\n        self.register_buffer('pos_enc', positional_encoding(n_patches + 1, 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.head = nn.Linear(embed_dim, 4)\n\n    def forward(self, x):\n        p = patchify(x, self.patch_size)\n        tok = self.embed(p)\n        cls = self.cls_token.expand(x.shape[0], -1, -1)\n        tok = torch.cat([cls, tok], dim=1) + self.pos_enc\n        out = self.encoder(tok)\n        return torch.sigmoid(self.head(out[:, 0]))  # regress from the [CLS] token's output\n\ndef iou_batch(pred, target):\n    px0, py0, pw, ph = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]\n    tx0, ty0, tw, th = target[:, 0], target[:, 1], target[:, 2], target[:, 3]\n    px1, py1, tx1, ty1 = px0 + pw, py0 + ph, tx0 + tw, ty0 + th\n    ix0, iy0 = torch.maximum(px0, tx0), torch.maximum(py0, ty0)\n    ix1, iy1 = torch.minimum(px1, tx1), torch.minimum(py1, ty1)\n    inter = (ix1 - ix0).clamp(min=0) * (iy1 - iy0).clamp(min=0)\n    union = pw * ph + tw * th - inter\n    return inter / union.clamp(min=1e-8)\n\nrng = np.random.default_rng(1)\nN = 400\nXtr, Btr = [], []\nfor _ in range(N):\n    img, box, _ = make_detection_example(rng, BASE_SHAPES)  # base classes only\n    Xtr.append(img); Btr.append(box)\nXtr, Btr = np.array(Xtr, dtype=np.float32), np.array(Btr, dtype=np.float32)\n\ntorch.manual_seed(0)\nbox_model = BoxRegressor()\nopt = torch.optim.Adam(box_model.parameters(), lr=0.001)\nXt, Bt = torch.tensor(Xtr).unsqueeze(1), torch.tensor(Btr)\nfor _ in range(400):\n    opt.zero_grad()\n    loss = F.mse_loss(box_model(Xt), Bt)\n    loss.backward()\n    opt.step()\n\n# test localization on ALL THREE shapes, including the never-boxed novel square\ntest_rng = np.random.default_rng(2)\nXte, Bte, labels_te = [], [], []\nfor _ in range(150):\n    img, box, lbl = make_detection_example(test_rng, SHAPES)\n    Xte.append(img); Bte.append(box); labels_te.append(lbl)\nXte, Bte = np.array(Xte, dtype=np.float32), np.array(Bte, dtype=np.float32)\n\nwith torch.no_grad():\n    pred_boxes = box_model(torch.tensor(Xte).unsqueeze(1))\nious = iou_batch(pred_boxes, torch.tensor(Bte)).numpy()\n\nfor shape in SHAPES:\n    mask = np.array([l == shape for l in labels_te])\n    print(f'{shape:>8} (base={shape in BASE_SHAPES}): mean box IoU = {ious[mask].mean():.3f}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "257a15a9",
   "source": "Localization transfers to the novel shape reasonably well — never having seen a labeled square box, the regressor still finds squares with a respectable IoU, though noticeably less precisely than the base classes it was actually trained on. Class-agnostic \"is there an object here\" is a more generic, more transferable skill than exact class identity.\n\n## Half 2: a CLIP-style classifier, trained on all three classes\n\nThis is the \"broad pretraining\" stand-in — a dual Transformer image/text encoder (Lesson {{lesson:vision_language_models}}'s pattern: a ViT-style image encoder, and a small self-attention text encoder over the caption tokens), contrastively trained on labeled crops of **all three** shapes, plus and circle and square. Real open-vocabulary detectors rely on exactly this asymmetry: the image-text encoder was pretrained on a vastly broader vocabulary (hundreds of millions of image-caption pairs) than any detection dataset's box annotations ever cover, so it already \"knows\" what a square looks like even though the detector's box-training data never had one.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "8c15146b",
   "source": "VOCAB = ['a', 'photo', 'of', 'plus', 'circle', 'square']\nW2ID = {w: i for i, w in enumerate(VOCAB)}\n\ndef caption_for(shape):\n    return [W2ID[w] for w in ['a', 'photo', 'of', shape]]\n\nclass ImageEncoder(nn.Module):\n    \"\"\"ViT-style: patchify, [CLS] token, Transformer encoder (Lesson 47's TinyViT pattern),\n    pooled via the [CLS] token to one embedding per crop.\"\"\"\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        n_patches = (img_size // patch_size) ** 2\n        patch_dim = 1 * patch_size * patch_size\n        self.embed = nn.Linear(patch_dim, embed_dim)\n        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)\n        self.register_buffer('pos_enc', positional_encoding(n_patches + 1, 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        p = patchify(x, self.patch_size)\n        tok = self.embed(p)\n        cls = self.cls_token.expand(x.shape[0], -1, -1)\n        tok = torch.cat([cls, tok], dim=1) + self.pos_enc\n        out = self.encoder(tok)\n        return self.proj(out[:, 0])\n\nclass TextEncoder(nn.Module):\n    \"\"\"A small self-attention text encoder over the caption tokens, mean-pooled to one\n    vector -- real CLIP's text tower is a Transformer too, not a mean-pooled embedding\n    lookup (the course's earlier, simpler stand-in).\"\"\"\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\ndef clip_loss(img_emb, txt_emb, temperature=0.1):\n    img_emb = F.normalize(img_emb, dim=1)\n    txt_emb = F.normalize(txt_emb, dim=1)\n    logits = img_emb @ txt_emb.T / temperature\n    targets = torch.arange(img_emb.shape[0])\n    return (F.cross_entropy(logits, targets) + F.cross_entropy(logits.T, targets)) / 2\n\n# train on CROPS, matching how the detector will feed it patches at inference time --\n# training on full uncropped images instead would create a train/test mismatch in object\n# scale. Each crop's box is also randomly jittered in position and size: a ViT has no\n# built-in translation invariance (Lesson 47's own finding), so without this augmentation\n# it overfits to perfectly-centered training crops and falls apart on the detector's own\n# slightly-imprecise predicted boxes later -- real CLIP-style training uses exactly this\n# kind of random-crop augmentation for the same robustness reason.\nclip_rng = np.random.default_rng(3)\nN_CLIP = 400\nX_clip, cap_clip = [], []\nfor _ in range(N_CLIP):\n    shape = clip_rng.choice(SHAPES)\n    cx, cy = clip_rng.integers(8, SIZE - 8), clip_rng.integers(8, SIZE - 8)\n    img = make_image(shape, cx, cy, r=6)\n    img = np.clip(img + clip_rng.normal(0, 0.05, img.shape), 0, 1).astype(np.float32)\n    jx, jy, jr = clip_rng.uniform(-2, 2), clip_rng.uniform(-2, 2), clip_rng.uniform(0.8, 1.2)\n    half = 6 * jr\n    box = np.array([cx - half + jx, cy - half + jy, 2 * half, 2 * half], dtype=np.float32) / SIZE\n    X_clip.append(crop_and_resize(img, box)); cap_clip.append(caption_for(shape))\nX_clip = np.array(X_clip, dtype=np.float32)\ncap_clip = np.array(cap_clip, dtype=np.int64)\n\ntorch.manual_seed(1)\nimg_enc = ImageEncoder()\ntxt_enc = TextEncoder(len(VOCAB))\nopt2 = torch.optim.Adam(list(img_enc.parameters()) + list(txt_enc.parameters()), lr=0.003)\ncaps_tensor = torch.tensor(cap_clip)\nn_clip = len(X_clip)\nfor epoch in range(800):\n    idx = np.random.default_rng(epoch).permutation(n_clip)[:64]\n    loss = clip_loss(img_enc(torch.tensor(X_clip[idx]).unsqueeze(1)), txt_enc(caps_tensor[idx]))\n    opt2.zero_grad()\n    loss.backward()\n    opt2.step()\n\nprint(f'CLIP training final loss: {loss.item():.3f}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "c6765387",
   "source": "## The full pipeline: detect, crop, classify by text match\n\nFor each test image: run the class-agnostic box regressor, crop the predicted region, embed it, and compare against text-prompt embeddings for all three shape names — including \"square,\" which the box regressor never trained on. To isolate *why* any errors happen, also evaluate the same classifier using the *ground-truth* box instead of the predicted one, which removes localization error from the picture entirely.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "846ff502",
   "source": "with torch.no_grad():\n    prompt_embs = {s: F.normalize(txt_enc(torch.tensor([caption_for(s)])), dim=1) for s in SHAPES}\n\ndef classify_crop(img, box):\n    with torch.no_grad():\n        crop = crop_and_resize(img, box)\n        emb = F.normalize(img_enc(torch.tensor(crop[None, None])), dim=1)\n        sims = {s: (emb @ prompt_embs[s].T).item() for s in SHAPES}\n    return max(sims, key=sims.get)\n\neval_rng = np.random.default_rng(2)\nn_correct_pred, n_correct_gt, n_total = {s: 0 for s in SHAPES}, {s: 0 for s in SHAPES}, {s: 0 for s in SHAPES}\nfor _ in range(150):\n    img, gt_box, true_shape = make_detection_example(eval_rng, SHAPES)\n    with torch.no_grad():\n        pred_box = box_model(torch.tensor(img[None, None])).squeeze(0).numpy()\n    pred_using_predicted_box = classify_crop(img, pred_box)\n    pred_using_gt_box = classify_crop(img, gt_box)\n    n_total[true_shape] += 1\n    n_correct_pred[true_shape] += (pred_using_predicted_box == true_shape)\n    n_correct_gt[true_shape] += (pred_using_gt_box == true_shape)\n\nprint(f'{\"class\":>8} {\"base?\":>7} {\"end-to-end (predicted box)\":>28} {\"classification only (true box)\":>32}')\nfor s in SHAPES:\n    acc_pred = n_correct_pred[s] / n_total[s]\n    acc_gt = n_correct_gt[s] / n_total[s]\n    print(f'{s:>8} {str(s in BASE_SHAPES):>7} {acc_pred:>27.1%} {acc_gt:>31.1%}')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "98caedd5",
   "source": "With the ground-truth box, classification hits 100% on all three classes — the CLIP-style classifier genuinely recognizes \"square\" by text prompt alone, despite the *detector* never once training on a square's box. That confirms the core open-vocabulary claim cleanly. But end-to-end, using the detector's own predicted box, square accuracy drops well behind the base classes — not because the classifier forgot what a square looks like, but because the class-agnostic box regressor's boxes for the novel shape are measurably less precise (recall the IoU gap from Half 1), and a distorted crop is harder to classify correctly even for a perfectly capable classifier. Novel-category *classification* and novel-category *localization* are genuinely separate problems with separate failure modes, and this experiment isolates exactly which one is responsible for any given end-to-end error — precisely the diagnostic real open-vocabulary detection papers report (base-class vs. novel-class AP, broken out separately) when describing where their systems still fall short.\n\n### Exercises\n\n1. Increase the training set size for `box_model` from 400 to 2000 base-class examples. Does more base-class training data improve novel-class (square) localization IoU, even though square itself is still never in the training labels?\n2. Add a fourth shape (e.g. a small triangle, following Lesson {{lesson:evaluating_a_classifier}}'s pattern) that is novel to *both* the box regressor and the CLIP classifier. Does the pipeline fail at the localization stage, the classification stage, or both — and how would you tell them apart using this lesson's ground-truth-box diagnostic?\n3. `crop_and_resize` always resizes to the same output size regardless of the predicted box's size. Real open-vocabulary detectors are sensitive to box-size errors for exactly this reason — a box that's the wrong *aspect ratio* distorts the crop before classification ever sees it. Measure the correlation between predicted-box IoU and classification correctness across the square test cases: do the worst-localized squares also tend to be the misclassified ones?",
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}