""" 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()