diff --git a/recognization/code/engine.py b/recognization/code/engine.py index 7da5f13..2d5c4e1 100644 --- a/recognization/code/engine.py +++ b/recognization/code/engine.py @@ -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: diff --git a/recognization/code/finetune_identity.py b/recognization/code/finetune_identity.py index ee6544b..cc7e519 100644 --- a/recognization/code/finetune_identity.py +++ b/recognization/code/finetune_identity.py @@ -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 /identity_head.pt for the gallery to load") args = parser.parse_args() + if args.stock: + train_from_stock(args) + return import json diff --git a/recognization/code/gallery.py b/recognization/code/gallery.py index f01a255..313538d 100644 --- a/recognization/code/gallery.py +++ b/recognization/code/gallery.py @@ -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) diff --git a/recognization/docs/EXTERNAL-EVAL.md b/recognization/docs/EXTERNAL-EVAL.md index 6693d2f..4bb0d9a 100644 --- a/recognization/docs/EXTERNAL-EVAL.md +++ b/recognization/docs/EXTERNAL-EVAL.md @@ -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 `/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.