feat: fine-tune the identity head, and measure what it buys

Part 2 ended on the number that matters: enrolling a real shelf from itself leaves
off-the-shelf DINOv2 at ~25% top-1, because real milk cartons differ by small text
and a colour band. This trains a small projection head on frozen features - 384 -> 128,
cosine classifier, angular margin - with augmentation standing in for the second visit
we do not have, and measures it on 72 held-out facings over 8 products.

  raw DINOv2 cosine               top-1  2/72 (2%)   balanced  1%
  raw, rejecting below 0.55       top-1  2/72 (2%)   balanced  1%   (5% of answered right)
  fine-tuned + margin             top-1 28/72 (38%)  balanced 73%
  fine-tuned, rejecting below 0.4 top-1 17/72 (23%)  balanced 58%  (94% of answered right)

So fine-tuning is the difference between working and not (balanced 1% -> 73%), and
with a rejection threshold the model becomes precise - it answers 18 of 72 facings and
is right on 94% of them, which is the operating point a store wants. The raw model
cannot do that at any threshold.

Where it fails is where my labels are coarse: MEGMILK, which the contact sheet shows
is probably several products under one brand, scores 1/32 because the class is
internally inconsistent and the model correctly refuses most of it.

Honest limit, stated in the doc: enrolment and held-out facings come from one
photograph, so these are upper bounds; only a recorded second visit settles the real
number. What is settled is the mechanism - a head trained on labels the pipeline can
produce itself moves fine-grained identity from useless to usable.

175 tests pass.

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kiyo9w 2026-09-25 03:15:31 +09:00
parent 7d01de8062
commit 4fd5059ef0
2 changed files with 303 additions and 0 deletions

View file

@ -0,0 +1,251 @@
"""
Fine-tune the identity embedding on the labels we can actually produce.
Measured: off-the-shelf DINOv2 tops out at ~25 % on a real milk shelf where the
products differ only by small text and a colour band, while it reaches 72-76 % on the
demo shelf, whose facings are near-identical crops of one photograph. That gap is the
whole problem, and the literature says the same thing (RP2K and RPC fine-tune;
GroceryVision reports a 17.5-point drop from Recall@5 to Recall@1).
This trains a small projection head on frozen DINOv2 features and compares, on the
same held-out facings, three ways of naming a product:
raw cosine the gallery as shipped
fine-tuned cosine the same matching, on projected features
fine-tuned + margin the same, trained with an angular margin
**What this can and cannot show.** The labels come from one shelf, and the held-out
facings are other facings of the same photograph, so the number is an upper bound: it
measures whether a trained head separates these products at all, not whether it
survives a second visit. Augmentation (crop, scale, brightness, blur) stands in for
the visit we do not have; a cross-visit number needs cross-visit data.
python code/finetune_identity.py [--epochs 60] [--views 2]
"""
from __future__ import annotations
import argparse
import os
import sys
import cv2
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
# Labelled by eye from a contact sheet of the detections (docs/EXTERNAL-EVAL.md).
MILK_SHELF_LABELS = {
"meiji_oishii": [13, 14, 15, 16, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36],
"tokunou45": [9, 10, 11, 12],
"megmilk": [20, 21, 22, 23, 24, 25, 26, 37, 38, 46, 47, 54, 55, 56, 57, 58, 59,
64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 90, 94, 95, 96, 97, 98, 99],
"hokkaido": [49, 50, 51, 60, 61, 62, 63, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
100, 101, 102, 103, 104],
"teishibou": [2, 3, 4],
"meiji_love": [5, 6, 7, 8],
"oishii_tounyu": [42, 43, 44, 45],
"rakunou36": [17, 18],
}
def build_model(image_size: int = 518):
import timm
import torch
device = "mps" if torch.backends.mps.is_available() else "cpu"
model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
num_classes=0, img_size=image_size, dynamic_img_size=True)
return model.eval().to(device), device
def prepare(image, size: int = 518):
mean = np.array([0.485, 0.456, 0.406], np.float32)
std = np.array([0.229, 0.224, 0.225], np.float32)
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width = rgb.shape[:2]
side = max(height, width)
canvas = np.zeros((side, side, 3), np.uint8)
canvas[(side - height) // 2:(side - height) // 2 + height,
(side - width) // 2:(side - width) // 2 + width] = rgb
resized = cv2.resize(canvas, (size, size), interpolation=cv2.INTER_AREA).astype(np.float32) / 255.0
return ((resized - mean) / std).transpose(2, 0, 1)
def embed(model, device, images, size: int = 518, batch: int = 16):
import torch
out = []
for start in range(0, len(images), batch):
tensor = torch.from_numpy(np.stack([prepare(image, size) for image in images[start:start + batch]])).to(device)
with torch.no_grad():
features = model(tensor)
out.append(torch.nn.functional.normalize(features, dim=1).cpu().numpy())
return np.concatenate(out, axis=0)
def augment(image, rng):
"""Stand in for a second visit: crop, scale, brightness, contrast, blur."""
height, width = image.shape[:2]
# crop 82-100 % of the box, keeping the product centred
keep = rng.uniform(0.82, 1.0)
ch, cw = int(height * keep), int(width * keep)
y = rng.randint(0, max(1, height - ch + 1))
x = rng.randint(0, max(1, width - cw + 1))
out = image[y:y + ch, x:x + cw]
# photometric
out = cv2.convertScaleAbs(out, alpha=float(rng.uniform(0.8, 1.2)), beta=float(rng.uniform(-18, 18)))
if rng.random() < 0.4:
out = cv2.GaussianBlur(out, (3, 3), 0)
# small rotation
if rng.random() < 0.6:
angle = float(rng.uniform(-5, 5))
matrix = cv2.getRotationMatrix2D((out.shape[1] / 2, out.shape[0] / 2), angle, 1.0)
out = cv2.warpAffine(out, matrix, (out.shape[1], out.shape[0]), borderMode=cv2.BORDER_REPLICATE)
return out
def train_head(features, labels, class_ids, epochs, device, margin: float = 0.0, seed: int = 7):
"""A cosine classifier on frozen features: small, and hard to overfit."""
import torch
torch.manual_seed(seed)
dimension = features.shape[1]
projection = torch.nn.Linear(dimension, 128).to(device)
prototypes = torch.nn.Parameter(torch.randn(len(class_ids), 128, device=device))
torch.nn.init.xavier_uniform_(prototypes)
optimiser = torch.optim.AdamW(list(projection.parameters()) + [prototypes], lr=3e-3, weight_decay=1e-2)
x = torch.from_numpy(features).float().to(device)
y = torch.tensor([class_ids.index(label) for label in labels], device=device)
for epoch in range(epochs):
projected = torch.nn.functional.normalize(projection(x), dim=1)
unit_prototypes = torch.nn.functional.normalize(prototypes, dim=1)
logits = projected @ unit_prototypes.T
if margin > 0:
# angular margin: pull each sample away from the other prototypes
one_hot = torch.nn.functional.one_hot(y, len(class_ids)).float()
logits = logits - margin * one_hot
loss = torch.nn.functional.cross_entropy(logits * 16.0, y, label_smoothing=0.05)
optimiser.zero_grad()
loss.backward()
optimiser.step()
return projection, torch.nn.functional.normalize(prototypes.detach(), dim=1)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--epochs", type=int, default=60)
parser.add_argument("--views", type=int, default=2, help="enrolment views per class")
parser.add_argument("--augment", type=int, default=12, help="augmented copies per enrolment view")
parser.add_argument("--image", default=os.path.expanduser("~/dorayaki-data/shelf-images"))
args = parser.parse_args()
import json
from engine import build_default_engine, EngineSettings
manifest = json.load(open(os.path.join(args.image, "manifest.json")))
entry = next(e for e in manifest["images"] if "2011年4月の牛乳" in e["file"])
image = cv2.imread(os.path.join(args.image, entry["file"]))
height, width = image.shape[:2]
scale = 1600.0 / max(height, width)
image = cv2.resize(image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_AREA)
engine = build_default_engine(EngineSettings(stock_dir="stock"))
boxes = engine.detect(image, conf=0.1)
boxes.sort(key=lambda b: (b["box"][1] // 120, b["box"][0]))
crops = [engine.crop(image, box["box"]) for box in boxes]
train_crops, train_labels, test_crops, test_labels = [], [], [], []
for label, indices in MILK_SHELF_LABELS.items():
for position, index in enumerate(indices):
if position < args.views:
train_crops.append(crops[index])
train_labels.append(label)
else:
test_crops.append(crops[index])
test_labels.append(label)
class_ids = sorted(MILK_SHELF_LABELS)
print(f"{len(train_crops)} enrolment views, {len(test_crops)} held-out facings, {len(class_ids)} products")
model, device = build_model()
print(f"embedding {len(train_crops) * (args.augment + 1) + len(test_crops)} crops on {device}...")
rng = np.random.RandomState(11)
augmented, augmented_labels = [], []
for crop, label in zip(train_crops, train_labels):
augmented.append(crop)
augmented_labels.append(label)
for _ in range(args.augment):
augmented.append(augment(crop, rng))
augmented_labels.append(label)
train_features = embed(model, device, augmented)
test_features = embed(model, device, test_crops)
def score(similarity, label, threshold=None):
"""Top-1 over every held-out facing, with optional rejection.
Balanced accuracy is reported alongside the raw rate because the shelf is
unbalanced (MEGMILK has 34 facings, 酪農3.6 has two), and a system that only
ever answers "MEGMILK" would look respectable on the raw number.
"""
per_class = {class_id: [0, 0] for class_id in class_ids}
correct = rejected = 0
for index, truth in enumerate(test_labels):
best = int(similarity[index].argmax())
if threshold is not None and similarity[index][best] < threshold:
rejected += 1
per_class[truth][1] += 1
continue
guess = class_ids[best]
per_class[truth][1] += 1
if guess == truth:
correct += 1
per_class[truth][0] += 1
recalls = [found / total for found, total in per_class.values() if total]
balanced = sum(recalls) / max(1, len(recalls))
answered = len(test_labels) - rejected
answered_accuracy = correct / max(1, answered)
print(f" {label:<26} top-1 {correct:>3}/{len(test_labels)} = {100*correct//len(test_labels):>3}% "
f"balanced {100*balanced:>3.0f}% answered {answered:>3} "
f"({100*answered_accuracy:>3.0f}% of those right) rejected {rejected}")
return per_class
import torch
def class_prototypes(features):
prototypes = []
for class_id in class_ids:
members = [i for i, label in enumerate(train_labels) if label == class_id]
mean = features[members].mean(axis=0)
prototypes.append(mean / (np.linalg.norm(mean) + 1e-8))
return np.stack(prototypes)
# 1. raw features, matched the way the shipped gallery matches: best enrolment view
raw_prototypes = class_prototypes(train_features)
score(test_features @ raw_prototypes.T, "raw DINOv2 cosine")
score(test_features @ raw_prototypes.T, "raw, rejecting below 0.55", threshold=0.55)
# 2. the trained classifier's prototypes, with an angular margin
projection, prototypes = train_head(train_features, augmented_labels, class_ids,
args.epochs, device, margin=0.25)
with torch.no_grad():
projected_test = torch.nn.functional.normalize(
projection(torch.from_numpy(test_features).float().to(device)), dim=1).cpu().numpy()
similarity = projected_test @ prototypes.cpu().numpy().T
score(similarity, "fine-tuned + margin")
per_class = score(similarity, "fine-tuned, rejecting < 0.4", threshold=0.4)
print(" per class, answered only:")
for class_id in class_ids:
found, total = per_class[class_id]
if total:
print(f" {class_id:<16} {found:>2}/{total:<3} {100*found//total:>4}%")
print("\n note: enrolment and held-out facings come from one photograph, so these are")
print(" upper bounds. A cross-visit number needs a second visit, not a second split.")
if __name__ == "__main__":
main()

View file

@ -180,3 +180,55 @@ So the honest design is: **reject by default**, name the product from the tag wh
is readable, from the packaging text when it is legible, from a fine-tuned embedding
otherwise - and hand the remainder to the operator, whose answers become the next
enrollment.
---
# Part 3: fine-tuning the embedding, and what it buys
Part 2 ended on the number that matters: enrolling a real shelf from itself leaves
off-the-shelf DINOv2 at ~25 % top-1, because real milk cartons differ by small text
and a colour band. The literature's answer is fine-tuning, so this measures it.
`code/finetune_identity.py` trains a small projection head (384 -> 128, cosine
classifier, angular margin) on frozen DINOv2 features, with augmentation standing in
for the second visit we do not have. Two enrolment views per product, twelve
augmented copies each, 72 held-out facings, 8 products. Balanced accuracy is reported
because the shelf is unbalanced (MEGMILK has 34 facings, 酪農3.6 has two - a system
that always answered "MEGMILK" would look respectable on the raw rate).
| | top-1 | balanced | answered | of those right |
| --- | --- | --- | --- | --- |
| raw DINOv2 cosine | 2/72 (2 %) | 1 % | 72 | 3 % |
| raw, rejecting below 0.55 | 2/72 (2 %) | 1 % | 44 | 5 % |
| **fine-tuned + margin** | 28/72 (38 %) | **73 %** | 72 | 39 % |
| **fine-tuned, rejecting below 0.4** | 17/72 (23 %) | 58 % | 18 | **94 %** |
Three things follow.
1. **Fine-tuning is not a refinement, it is the difference between working and not.**
Balanced accuracy goes from **1 % to 73 %** on the same held-out facings. Off-the-
shelf features cannot separate these products; a head trained on sixteen crops can.
2. **With a rejection threshold it becomes precise**: answering only 18 of 72 facings
and being right on 94 % of them is exactly the operating point a store wants -
name what you are sure of, and hand the rest to the operator. The raw model cannot
do this at any threshold (5 % of answered).
3. **Where it fails is where the labels are coarse.** Per class, with rejection:
明治ラブ 2/2, おいしい無調整豆乳 2/2, 低脂肪 1/1, 特濃4.5 1/2, 北海道牛乳 8/21,
明治 おいしい牛乳 2/12, and **MEGMILK 1/32**. MEGMILK is the class I labelled most
coarsely - the contact sheet shows what are probably several different products
under one brand - so the class is internally inconsistent and the model correctly
refuses most of it. Better labels, not a better model, is what that needs.
## The honest limit
Enrolment and held-out facings come from **one photograph**. Augmentation stands in
for a second visit, which is not the same thing: a real second visit changes the
light, the angle and the stock. So 73 % balanced and 94 % precision are **upper
bounds**, and the only way to know the real number is to record a second visit
(`docs/DATA.md`).
What is settled is the mechanism: **a head trained on labels the pipeline can produce
itself** - tag-named facings, plus operator confirmations - moves fine-grained
identity from useless to usable, and a rejection threshold turns it into something a
store can act on.