404 lines
17 KiB
Python
404 lines
17 KiB
Python
"""
|
|
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 enrollment_report(stock: str, csv_name: str = "stock.csv") -> dict:
|
|
"""Catalog versus enrolled views. Does not load a model or write files.
|
|
|
|
A SKU with no readable view cannot be recognized. One view is reported, not
|
|
treated as enough: the measured knee is two views. Duplicate CSV rows and a
|
|
head whose class list disagrees with the catalog are blockers.
|
|
"""
|
|
import csv
|
|
import re
|
|
|
|
from gallery import ProductGallery
|
|
|
|
names = []
|
|
try:
|
|
names = os.listdir(stock)
|
|
except OSError:
|
|
names = []
|
|
images = [name for name in names if os.path.splitext(name)[1].lower() in (".jpeg", ".jpg", ".png")]
|
|
rows = []
|
|
csv_path = os.path.join(stock, csv_name)
|
|
if os.path.exists(csv_path):
|
|
with open(csv_path, encoding="utf-8") as handle:
|
|
rows = list(csv.DictReader(handle))
|
|
sku_ids = []
|
|
seen = set()
|
|
duplicate_rows = []
|
|
for row in rows:
|
|
sku_id = (row.get("SKU_ID") or "").strip()
|
|
if not sku_id:
|
|
continue
|
|
if sku_id in seen:
|
|
duplicate_rows.append(sku_id)
|
|
seen.add(sku_id)
|
|
sku_ids.append(sku_id)
|
|
gallery = ProductGallery(stock, sku_ids or [])
|
|
views = {sku_id: gallery._views_for(sku_id) for sku_id in dict.fromkeys(sku_ids)}
|
|
missing = [sku_id for sku_id, paths in views.items() if not paths]
|
|
single = [sku_id for sku_id, paths in views.items() if len(paths) == 1]
|
|
claimed = {os.path.basename(path) for paths in views.values() for path in paths}
|
|
orphans = sorted(name for name in images if name not in claimed)
|
|
# ponytail: prefix collision check. A view named sku_001 must not also match sku_00.
|
|
collisions = []
|
|
for sku_id, paths in views.items():
|
|
bare = [os.path.basename(path) for path in paths if os.path.splitext(os.path.basename(path))[0] != sku_id
|
|
and not re.fullmatch(re.escape(sku_id) + r"[_-]\d+", os.path.splitext(os.path.basename(path))[0])]
|
|
if bare:
|
|
collisions.append(sku_id)
|
|
blockers = []
|
|
if duplicate_rows:
|
|
blockers.append("duplicate_csv_sku")
|
|
if missing:
|
|
blockers.append("sku_without_view")
|
|
if collisions:
|
|
blockers.append("view_prefix_collision")
|
|
head_path = os.path.join(stock, "identity_head.pt")
|
|
head_note = None
|
|
if os.path.exists(head_path):
|
|
head_note = "present_not_opened"
|
|
# Do not torch.load here: the head can be large, and this report must stay cheap.
|
|
# Class-list drift is reported when the caller passes class_ids.
|
|
return {
|
|
"stock": stock,
|
|
"sku_ids": list(dict.fromkeys(sku_ids)),
|
|
"images": len(images),
|
|
"missing": missing,
|
|
"single_view": single,
|
|
"orphans": orphans,
|
|
"duplicate_rows": duplicate_rows,
|
|
"collisions": collisions,
|
|
"blockers": blockers,
|
|
"head": head_note,
|
|
}
|
|
|
|
|
|
def enrollment_report_text(report: dict) -> str:
|
|
return (
|
|
f"{report['stock']}: {len(report['sku_ids'])} SKUs, {report['images']} images, "
|
|
f"missing={len(report['missing'])}, single_view={len(report['single_view'])}, "
|
|
f"orphans={len(report['orphans'])}, blockers={report['blockers'] or 'none'}"
|
|
)
|
|
|
|
|
|
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 train_from_stock(args) -> None:
|
|
"""Train the head on a store's own enrolled views, for the gallery to load.
|
|
|
|
The head belongs to the store: the views in its stock directory are its products,
|
|
and a head trained on them is what turns the gallery from a lookup that cannot
|
|
tell two milk cartons apart into one that can. One to four views per SKU is
|
|
enough, with augmentation standing in for the rest.
|
|
"""
|
|
import glob
|
|
import json
|
|
|
|
import torch
|
|
|
|
from gallery import ProductGallery
|
|
|
|
stock = args.stock
|
|
report = enrollment_report(stock)
|
|
print(enrollment_report_text(report))
|
|
if report["blockers"]:
|
|
print(f"{stock}: not training; enrollment blockers: {', '.join(report['blockers'])}")
|
|
return
|
|
gallery = ProductGallery(stock, report["sku_ids"])
|
|
views = {sku: gallery._views_for(sku) for sku in report["sku_ids"]}
|
|
views = {sku: paths for sku, paths in views.items() if paths}
|
|
if len(views) < 2:
|
|
print(f"{stock}: needs at least two products with views, found {len(views)}")
|
|
return
|
|
|
|
crops, labels = [], []
|
|
for sku, paths in sorted(views.items()):
|
|
for path in paths:
|
|
image = cv2.imread(path)
|
|
if image is not None and image.size:
|
|
crops.append(image)
|
|
labels.append(sku)
|
|
class_ids = sorted(views)
|
|
print(f"{stock}: {len(crops)} views over {len(class_ids)} products")
|
|
|
|
model, device = build_model()
|
|
rng = np.random.RandomState(11)
|
|
augmented, augmented_labels = [], []
|
|
for crop, label in zip(crops, labels):
|
|
augmented.append(crop)
|
|
augmented_labels.append(label)
|
|
for _ in range(args.augment):
|
|
augmented.append(augment(crop, rng))
|
|
augmented_labels.append(label)
|
|
print(f"embedding {len(augmented)} crops on {device}...")
|
|
features = embed(model, device, augmented)
|
|
|
|
projection, prototypes = train_head(features, augmented_labels, class_ids,
|
|
args.epochs, device, margin=0.25)
|
|
payload = {
|
|
"class_ids": class_ids,
|
|
"image_size": 518,
|
|
"base_model": "vit_small_patch14_dinov2.lvd142m",
|
|
"projection": {key: value.cpu() for key, value in projection.state_dict().items()},
|
|
"prototypes": prototypes.cpu(),
|
|
"views": {sku: len(paths) for sku, paths in sorted(views.items())},
|
|
}
|
|
path = os.path.join(stock, "identity_head.pt")
|
|
torch.save(payload, path)
|
|
print(f"wrote {path} ({len(class_ids)} products, {len(crops)} views)")
|
|
|
|
|
|
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"))
|
|
parser.add_argument("--stock", default=None,
|
|
help="train on a store's enrolled views instead of the milk shelf, "
|
|
"and write <stock>/identity_head.pt for the gallery to load")
|
|
args = parser.parse_args()
|
|
if args.stock:
|
|
train_from_stock(args)
|
|
return
|
|
|
|
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()
|