vnguyen/product-recognition-demo
1
0
Fork 1
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

251
recognization/code/finetune_identity.py Normal file
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()