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

feat: a store's fine-tuned head, and the two-stage gate it needs

Browse source 
The head is now part of a store rather than a script: `finetune_identity.py --stock
stock-lopia` trains on that store's enrolled views (27 views over 7 products, 12
augmented copies each) and writes `<stock>/identity_head.pt`, which ProductGallery
loads automatically.

The first attempt exposed a trap worth recording. A discriminative head has no "none
of the above": trained with an angular margin to separate seven products, it maps
everything confidently onto those seven, including cartons it has never seen.

  raw features, threshold 0.55   home shelf 78% accepted   foreign shelves  5% accepted
  trained head, threshold 0.40   home shelf 58% accepted   foreign shelves 44% accepted

Rejection - the gallery's whole point - disappears. The fix needs no retraining: the
two mechanisms fail in opposite directions, so raw features decide *whether* to answer
(they generalise) and the head decides *what* to answer (it discriminates). Measured:
47% accepted on the home shelf, 5% on foreign ones. Lower coverage than raw-only,
which is the honest trade.

The head also needs its own threshold (0.40 against the raw 0.55): a trained head's
similarity scale is its own, and reusing the raw line rejected almost everything it
said.

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:21:22 +09:00
parent 4fd5059ef0
commit 2bf916f2aa
4 changed files with 192 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

2
recognization/code/engine.py
View file

@ -93,6 +93,7 @@ class EngineSettings:
# Cosine similarity below which the gallery calls a facing unknown. Measured:
# same-product pairs sit above 0.6, different products below 0.5.
gallery_threshold: float = 0.55
gallery_head_threshold: float = 0.40
gallery_image_size: int = 518
imgsz: int = 640
# 0 -> max(2, cpu_count - 2). Measured default below (see build_default_engine).
@ -430,6 +431,7 @@ class RecognitionEngine:
self.settings.stock_dir,
[product["sku_id"] for product in self.catalog()],
threshold=self.settings.gallery_threshold,
head_threshold=self.settings.gallery_head_threshold,
image_size=self.settings.gallery_image_size,
)
try:

71
recognization/code/finetune_identity.py
View file

@ -133,13 +133,84 @@ def train_head(features, labels, class_ids, epochs, device, margin: float = 0.0,
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
gallery = ProductGallery(stock, [])
# the SKUs are whatever the stock directory holds
skus = sorted({os.path.splitext(name)[0].rsplit("_", 1)[0] if os.path.splitext(name)[0].rsplit("_", 1)[-1].isdigit()
else os.path.splitext(name)[0]
for name in os.listdir(stock)
if os.path.splitext(name)[1].lower() in (".jpeg", ".jpg", ".png")})
views = {sku: gallery._views_for(sku) for sku in skus}
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

70
recognization/code/gallery.py
View file

@ -34,6 +34,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
# from the measured distribution: same-product pairs sit above 0.6, different
# products below 0.5, and the gap is where a rejection belongs.
DEFAULT_THRESHOLD = 0.55
# The same line for a fine-tuned head, whose similarity scale is its own.
DEFAULT_HEAD_THRESHOLD = 0.40
@dataclass
@ -66,19 +68,50 @@ class ProductGallery:
device: Optional[str] = None,
image_size: int = 518,
threshold: float = DEFAULT_THRESHOLD,
head_threshold: float = DEFAULT_HEAD_THRESHOLD,
) -> None:
self.stock_dir = stock_dir
self.sku_ids = list(sku_ids)
self.image_size = int(image_size)
self.threshold = float(threshold)
# A trained head produces a different similarity scale than raw features, so
# it needs its own line. Measured on a real milk shelf: the head answers 18 of
# 72 facings above 0.4 and is right on 94 % of them, while the raw threshold
# (0.55) rejects almost everything the head says.
self.head_threshold = float(head_threshold)
self.device = device
self._model = None
self._lock = threading.Lock()
self._entries: List[GalleryEntry] = []
self._matrix = None
self._available = False
# A fine-tuned projection head, if the store has one (see
# finetune_identity.py). Measured on a real milk shelf: raw features put 1 %
# balanced accuracy on fine-grained products, a trained head 73 %.
self._head = None
self._head_path = os.path.join(stock_dir, "identity_head.pt")
# ------------------------------------------------------------------ loading
def load_head(self) -> bool:
"""Load the store's fine-tuned head, if it has one."""
import torch
if not os.path.exists(self._head_path):
return False
try:
payload = torch.load(self._head_path, map_location="cpu", weights_only=False)
except Exception:
return False
projection = torch.nn.Linear(384, 128)
projection.load_state_dict(payload["projection"])
projection.eval()
self._head = {
"projection": projection,
"prototypes": torch.nn.functional.normalize(payload["prototypes"].float(), dim=1),
"class_ids": list(payload["class_ids"]),
}
return True
def _load_model(self) -> Any:
if self._model is not None:
return self._model
@ -156,6 +189,7 @@ class ProductGallery:
]
self._matrix = features
self._available = True
self.load_head()
return True
# ------------------------------------------------------------------ querying
@ -171,13 +205,42 @@ class ProductGallery:
return features[0]
def match(self, crop: Any, threshold: Optional[float] = None) -> GalleryMatch:
"""Nearest enrolled view, with a rejection when nothing is close enough."""
"""Nearest enrolled product, with a rejection when nothing is close enough."""
if not self._available or self._matrix is None:
return GalleryMatch(sku_id=None, similarity=0.0, rejected=True)
import numpy as np
import torch
limit = self.threshold if threshold is None else float(threshold)
limit = threshold
if limit is None:
limit = self.head_threshold if self._head is not None else self.threshold
limit = float(limit)
embedding = self.embed(crop)
if self._head is not None:
# Two stages, because the two mechanisms fail in opposite directions.
# Measured: raw features reject strangers well (5 % false accepts) but
# rank fine-grained products badly (2 % top-1); a trained head ranks them
# well (73 % balanced) but is confidently wrong on strangers (44 % false
# accepts), because a discriminative head has no "none of the above".
# So the raw features decide *whether* to answer, and the head decides
# *what* to answer.
raw_similarity = self._matrix @ embedding
best_view = float(raw_similarity.max()) if len(raw_similarity) else 0.0
if best_view < self.threshold:
return GalleryMatch(sku_id=None, similarity=best_view, rejected=True)
with torch.no_grad():
projected = torch.nn.functional.normalize(
self._head["projection"](torch.from_numpy(embedding).float()), dim=0
)
similarity = (self._head["prototypes"] @ projected).numpy()
ranked = sorted(zip(self._head["class_ids"], similarity), key=lambda kv: kv[1], reverse=True)
top_sku, top_score = ranked[0]
runner_sku, runner_score = ranked[1] if len(ranked) > 1 else (None, 0.0)
if float(top_score) < limit:
return GalleryMatch(sku_id=None, similarity=float(top_score), runner_up=top_sku,
runner_up_similarity=float(top_score), rejected=True)
return GalleryMatch(sku_id=top_sku, similarity=float(top_score), runner_up=runner_sku,
runner_up_similarity=float(runner_score))
similarity = self._matrix @ embedding
# Several views per SKU: a SKU's score is its best view.
best_by_sku: Dict[str, float] = {}
@ -196,6 +259,9 @@ class ProductGallery:
return GalleryMatch(sku_id=top_sku, similarity=top_score, runner_up=runner_sku,
runner_up_similarity=runner_score)
def head_loaded(self) -> bool:
return self._head is not None
def view_count(self) -> int:
return len(self._entries)

51
recognization/docs/EXTERNAL-EVAL.md
View file

@ -232,3 +232,54 @@ What is settled is the mechanism: **a head trained on labels the pipeline can pr
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.
---
# Part 4: making it usable, and the trap that came with it
The head is now part of a store: `python code/finetune_identity.py --stock stock-lopia`
trains on that store's enrolled views and writes `<stock>/identity_head.pt`, which
`ProductGallery` loads automatically. The first attempt exposed a trap worth writing
down.
## A trained head destroys rejection
Measured on the same two sets (the home shelf it was enrolled from, and 20 shelves
whose products are absent):
| | home shelf accepted | foreign shelves accepted |
| --- | --- | --- |
| raw features, threshold 0.55 | 78 % | **5 %** |
| trained head, threshold 0.40 | 58 % | **44 %** |
A discriminative head has no "none of the above": trained with an angular margin to
separate seven products, it maps *everything* confidently onto its seven prototypes,
including cartons it has never seen. Rejection, which was the gallery's whole point,
disappears.
## The fix is two stages, and it needs no retraining
The two mechanisms fail in opposite directions, so each does what it is good at:
* **raw features decide whether to answer** (they generalise: 5 % false accepts on
strangers),
* **the head decides what to answer** (it discriminates: 73 % balanced on
fine-grained products against 1 % raw).
| | home shelf accepted | foreign shelves accepted |
| --- | --- | --- |
| **raw gate, then head ranking** | 47 % | **5 %** |
The home-shelf acceptance is lower than raw-only because the gate is strict, which is
the honest trade: fewer claims, and the claims are worth more.
## Still open
* A joint accuracy number on the LOPIA store needs held-out labels there; its
enrolment views came from the facings the tags named, so scoring on those would be
scoring on the training set.
* Training the head *with* out-of-distribution negatives (crops from other stores)
should let it keep rejecting on its own, which would remove the need for the gate.
* Everything above is still one photograph per store. The second visit remains the
measurement that settles it.