Lesson 56: Open-Vocabulary Detection

Lesson 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 55, 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.

In [1]:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import matplotlib.patches as patches

Base classes vs. a novel class

Three 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.

In [2]:
SIZE = 32
SHAPES = ['plus', 'circle', 'square']
BASE_SHAPES = ['plus', 'circle']
NOVEL_SHAPE = 'square'

def make_image(shape_type, cx, cy, size=SIZE, r=6):
    img = np.zeros((size, size), dtype=np.float32)
    if shape_type == 'plus':
        img[cy-1:cy+2, cx-r:cx+r+1] = 1.0
        img[cy-r:cy+r+1, cx-1:cx+2] = 1.0
    elif shape_type == 'circle':
        yy, xx = np.mgrid[0:size, 0:size]
        img[((xx-cx)**2 + (yy-cy)**2) <= r**2] = 1.0
    else:  # square
        img[cy-r:cy+r+1, cx-r:cx+r+1] = 1.0
    return img

def make_detection_example(rng, shapes, size=SIZE, r=6):
    shape_type = rng.choice(shapes)
    cx, cy = rng.integers(r + 2, size - r - 2), rng.integers(r + 2, size - r - 2)
    img = make_image(shape_type, cx, cy)
    img = np.clip(img + rng.normal(0, 0.05, img.shape), 0, 1).astype(np.float32)
    box = np.array([cx - r, cy - r, 2 * r, 2 * r], dtype=np.float32) / size
    return img, box, shape_type

def crop_and_resize(img, box, out_size=SIZE):
    x0, y0, w, h = (box * SIZE).astype(int)
    x0, y0 = max(0, x0), max(0, y0)
    x1, y1 = min(SIZE, x0 + max(w, 1)), min(SIZE, y0 + max(h, 1))
    crop = img[y0:y1, x0:x1]
    if crop.size == 0:
        return np.zeros((out_size, out_size), dtype=np.float32)
    resized = F.interpolate(torch.tensor(crop[None, None]), size=(out_size, out_size),
                             mode='bilinear', align_corners=False)
    return resized[0, 0].numpy()

fig, axes = plt.subplots(1, 3, figsize=(7, 2.5))
demo_rng = np.random.default_rng(42)
for ax, shape in zip(axes, SHAPES):
    im, box, _ = make_detection_example(demo_rng, [shape])
    ax.imshow(im, cmap='gray')
    x0, y0, w, h = box * SIZE
    ax.add_patch(patches.Rectangle((x0, y0), w, h, edgecolor='lime', facecolor='none', linewidth=2))
    ax.set_title(f'{shape}  (base={shape in BASE_SHAPES})', fontsize=9)
    ax.axis('off')
plt.show()
No description has been provided for this image

Half 1: a class-agnostic box regressor

A ViT-style box regressor: Lesson 47'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 41's box regression: never a class label, only "where is the object," trained exclusively on plus and circle scenes.

In [3]:
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 BoxRegressor(nn.Module):
    def __init__(self, img_size=SIZE, patch_size=4, embed_dim=32, n_heads=4, n_layers=2):
        super().__init__()
        self.patch_size = patch_size
        n_patches = (img_size // patch_size) ** 2
        patch_dim = 1 * patch_size * patch_size
        self.embed = nn.Linear(patch_dim, embed_dim)
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
        self.register_buffer('pos_enc', positional_encoding(n_patches + 1, 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.head = nn.Linear(embed_dim, 4)

    def forward(self, x):
        p = patchify(x, self.patch_size)
        tok = self.embed(p)
        cls = self.cls_token.expand(x.shape[0], -1, -1)
        tok = torch.cat([cls, tok], dim=1) + self.pos_enc
        out = self.encoder(tok)
        return torch.sigmoid(self.head(out[:, 0]))  # regress from the [CLS] token's output

def iou_batch(pred, target):
    px0, py0, pw, ph = pred[:, 0], pred[:, 1], pred[:, 2], pred[:, 3]
    tx0, ty0, tw, th = target[:, 0], target[:, 1], target[:, 2], target[:, 3]
    px1, py1, tx1, ty1 = px0 + pw, py0 + ph, tx0 + tw, ty0 + th
    ix0, iy0 = torch.maximum(px0, tx0), torch.maximum(py0, ty0)
    ix1, iy1 = torch.minimum(px1, tx1), torch.minimum(py1, ty1)
    inter = (ix1 - ix0).clamp(min=0) * (iy1 - iy0).clamp(min=0)
    union = pw * ph + tw * th - inter
    return inter / union.clamp(min=1e-8)

rng = np.random.default_rng(1)
N = 400
Xtr, Btr = [], []
for _ in range(N):
    img, box, _ = make_detection_example(rng, BASE_SHAPES)  # base classes only
    Xtr.append(img); Btr.append(box)
Xtr, Btr = np.array(Xtr, dtype=np.float32), np.array(Btr, dtype=np.float32)

torch.manual_seed(0)
box_model = BoxRegressor()
opt = torch.optim.Adam(box_model.parameters(), lr=0.001)
Xt, Bt = torch.tensor(Xtr).unsqueeze(1), torch.tensor(Btr)
for _ in range(400):
    opt.zero_grad()
    loss = F.mse_loss(box_model(Xt), Bt)
    loss.backward()
    opt.step()

# test localization on ALL THREE shapes, including the never-boxed novel square
test_rng = np.random.default_rng(2)
Xte, Bte, labels_te = [], [], []
for _ in range(150):
    img, box, lbl = make_detection_example(test_rng, SHAPES)
    Xte.append(img); Bte.append(box); labels_te.append(lbl)
Xte, Bte = np.array(Xte, dtype=np.float32), np.array(Bte, dtype=np.float32)

with torch.no_grad():
    pred_boxes = box_model(torch.tensor(Xte).unsqueeze(1))
ious = iou_batch(pred_boxes, torch.tensor(Bte)).numpy()

for shape in SHAPES:
    mask = np.array([l == shape for l in labels_te])
    print(f'{shape:>8} (base={shape in BASE_SHAPES}): mean box IoU = {ious[mask].mean():.3f}')
    plus (base=True): mean box IoU = 0.874
  circle (base=True): mean box IoU = 0.915
  square (base=False): mean box IoU = 0.734

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.

Half 2: a CLIP-style classifier, trained on all three classes

This is the "broad pretraining" stand-in — a dual Transformer image/text encoder (Lesson 55'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.

In [4]:
VOCAB = ['a', 'photo', 'of', 'plus', 'circle', 'square']
W2ID = {w: i for i, w in enumerate(VOCAB)}

def caption_for(shape):
    return [W2ID[w] for w in ['a', 'photo', 'of', shape]]

class ImageEncoder(nn.Module):
    """ViT-style: patchify, [CLS] token, Transformer encoder (Lesson 47's TinyViT pattern),
    pooled via the [CLS] token to one embedding per crop."""
    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
        n_patches = (img_size // patch_size) ** 2
        patch_dim = 1 * patch_size * patch_size
        self.embed = nn.Linear(patch_dim, embed_dim)
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
        self.register_buffer('pos_enc', positional_encoding(n_patches + 1, 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):
        p = patchify(x, self.patch_size)
        tok = self.embed(p)
        cls = self.cls_token.expand(x.shape[0], -1, -1)
        tok = torch.cat([cls, tok], dim=1) + self.pos_enc
        out = self.encoder(tok)
        return self.proj(out[:, 0])

class TextEncoder(nn.Module):
    """A small self-attention text encoder over the caption tokens, mean-pooled to one
    vector -- real CLIP's text tower is a Transformer too, not a mean-pooled embedding
    lookup (the course's earlier, simpler stand-in)."""
    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 clip_loss(img_emb, txt_emb, temperature=0.1):
    img_emb = F.normalize(img_emb, dim=1)
    txt_emb = F.normalize(txt_emb, dim=1)
    logits = img_emb @ txt_emb.T / temperature
    targets = torch.arange(img_emb.shape[0])
    return (F.cross_entropy(logits, targets) + F.cross_entropy(logits.T, targets)) / 2

# train on CROPS, matching how the detector will feed it patches at inference time --
# training on full uncropped images instead would create a train/test mismatch in object
# scale. Each crop's box is also randomly jittered in position and size: a ViT has no
# built-in translation invariance (Lesson 47's own finding), so without this augmentation
# it overfits to perfectly-centered training crops and falls apart on the detector's own
# slightly-imprecise predicted boxes later -- real CLIP-style training uses exactly this
# kind of random-crop augmentation for the same robustness reason.
clip_rng = np.random.default_rng(3)
N_CLIP = 400
X_clip, cap_clip = [], []
for _ in range(N_CLIP):
    shape = clip_rng.choice(SHAPES)
    cx, cy = clip_rng.integers(8, SIZE - 8), clip_rng.integers(8, SIZE - 8)
    img = make_image(shape, cx, cy, r=6)
    img = np.clip(img + clip_rng.normal(0, 0.05, img.shape), 0, 1).astype(np.float32)
    jx, jy, jr = clip_rng.uniform(-2, 2), clip_rng.uniform(-2, 2), clip_rng.uniform(0.8, 1.2)
    half = 6 * jr
    box = np.array([cx - half + jx, cy - half + jy, 2 * half, 2 * half], dtype=np.float32) / SIZE
    X_clip.append(crop_and_resize(img, box)); cap_clip.append(caption_for(shape))
X_clip = np.array(X_clip, dtype=np.float32)
cap_clip = np.array(cap_clip, dtype=np.int64)

torch.manual_seed(1)
img_enc = ImageEncoder()
txt_enc = TextEncoder(len(VOCAB))
opt2 = torch.optim.Adam(list(img_enc.parameters()) + list(txt_enc.parameters()), lr=0.003)
caps_tensor = torch.tensor(cap_clip)
n_clip = len(X_clip)
for epoch in range(800):
    idx = np.random.default_rng(epoch).permutation(n_clip)[:64]
    loss = clip_loss(img_enc(torch.tensor(X_clip[idx]).unsqueeze(1)), txt_enc(caps_tensor[idx]))
    opt2.zero_grad()
    loss.backward()
    opt2.step()

print(f'CLIP training final loss: {loss.item():.3f}')
CLIP training final loss: 3.071

The full pipeline: detect, crop, classify by text match

For 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.

In [5]:
with torch.no_grad():
    prompt_embs = {s: F.normalize(txt_enc(torch.tensor([caption_for(s)])), dim=1) for s in SHAPES}

def classify_crop(img, box):
    with torch.no_grad():
        crop = crop_and_resize(img, box)
        emb = F.normalize(img_enc(torch.tensor(crop[None, None])), dim=1)
        sims = {s: (emb @ prompt_embs[s].T).item() for s in SHAPES}
    return max(sims, key=sims.get)

eval_rng = np.random.default_rng(2)
n_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}
for _ in range(150):
    img, gt_box, true_shape = make_detection_example(eval_rng, SHAPES)
    with torch.no_grad():
        pred_box = box_model(torch.tensor(img[None, None])).squeeze(0).numpy()
    pred_using_predicted_box = classify_crop(img, pred_box)
    pred_using_gt_box = classify_crop(img, gt_box)
    n_total[true_shape] += 1
    n_correct_pred[true_shape] += (pred_using_predicted_box == true_shape)
    n_correct_gt[true_shape] += (pred_using_gt_box == true_shape)

print(f'{"class":>8} {"base?":>7} {"end-to-end (predicted box)":>28} {"classification only (true box)":>32}')
for s in SHAPES:
    acc_pred = n_correct_pred[s] / n_total[s]
    acc_gt = n_correct_gt[s] / n_total[s]
    print(f'{s:>8} {str(s in BASE_SHAPES):>7} {acc_pred:>27.1%} {acc_gt:>31.1%}')
   class   base?   end-to-end (predicted box)   classification only (true box)
    plus    True                      100.0%                          100.0%
  circle    True                      100.0%                          100.0%
  square   False                       97.9%                          100.0%

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.

Exercises

  1. 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?
  2. 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?
  3. 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?