{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 52: Open-Vocabulary Detection\n",
    "\n",
    "Lesson 40's detector predicts one of a small, fixed set of classes baked in at training time \u2014 exactly Lesson 37'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 49, 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 \u2014 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 40's box regression, unchanged, except the target is never a class label \u2014 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 \u2014 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 \u2014 a dual image/text encoder (Lesson 49), 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 \u2014 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 \u2014 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 \u2014 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 \u2014 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### Exercise\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 37'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 \u2014 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 \u2014 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
}