{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Lesson 55: Open-Vocabulary Detection\n\nLesson 41's detector predicts one of a small, fixed set of classes baked in at training time — exactly Lesson 36'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 52, 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\nLesson 41's box regression, unchanged, except the target is never a class label — only \"where is the object,\" trained exclusively on plus and circle scenes.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "4d0507fe",
   "source": "class BoxRegressor(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1))\n        self.fc = nn.Linear(32, 4)\n\n    def forward(self, x):\n        return torch.sigmoid(self.fc(self.conv(x).flatten(1)))\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.005)\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 image/text encoder (Lesson 52), 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    def __init__(self, out_dim=16):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(1, 16, 5, padding=2), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(16, 32, 5, padding=2), nn.ReLU(), nn.AdaptiveMaxPool2d(1))\n        self.proj = nn.Linear(32, out_dim)\n\n    def forward(self, x):\n        return self.proj(self.conv(x).flatten(1))\n\nclass TextEncoder(nn.Module):\n    def __init__(self, vocab_size, out_dim=16, embed_dim=8):\n        super().__init__()\n        self.embed = nn.Embedding(vocab_size, embed_dim)\n        self.proj = nn.Linear(embed_dim, out_dim)\n\n    def forward(self, ids):\n        return self.proj(self.embed(ids).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 scale\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    box = np.array([cx - 6, cy - 6, 12, 12], 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.01)\ncaps_tensor = torch.tensor(cap_clip)\nn_clip = len(X_clip)\nfor epoch in range(300):\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 36'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
}