fix: constrain trained head labels to enrolled catalog

recognization/code/finetune_identity.py

@@ -228,7 +228,7 @@ def train_from_stock(args) -> None:
 
     import torch
 
-    from gallery import ProductGallery
+    from gallery import DINOV2_MODEL, ProductGallery, catalog_digest, head_class_set
 
     stock = args.stock
     report = enrollment_report(stock)
@@ -267,10 +267,15 @@ def train_from_stock(args) -> None:
 
     projection, prototypes = train_head(features, augmented_labels, class_ids,
                                         args.epochs, device, margin=0.25)
+    digest = catalog_digest(stock, class_ids)
+    if head_class_set(class_ids) is None or digest is None:
+        print(f"{stock}: refusing to write a head; class ids are empty or duplicated")
+        return
     payload = {
         "class_ids": class_ids,
+        "catalog_digest": digest,
         "image_size": 518,
-        "base_model": "vit_small_patch14_dinov2.lvd142m",
+        "base_model": DINOV2_MODEL,
         "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())},

recognization/code/gallery.py

@@ -25,6 +25,7 @@ on the GPU. That is why this can sit in the live loop at all.
 
 from __future__ import annotations
 
+import hashlib
 import os
 import re
 import threading
@@ -131,6 +132,39 @@ class GalleryMatch:
         return self.similarity - self.runner_up_similarity
 
 
+def head_class_set(sku_ids: Sequence[str]) -> Optional[List[str]]:
+    """Catalog order, or None when an id is empty or repeated."""
+    cleaned = [str(sku_id).strip() for sku_id in sku_ids]
+    if not cleaned or any(not sku_id for sku_id in cleaned):
+        return None
+    if len(set(cleaned)) != len(cleaned):
+        return None
+    return cleaned
+
+
+def catalog_digest(stock_dir: str, sku_ids: Sequence[str], csv_name: str = "stock.csv") -> Optional[str]:
+    """Pin for identity_head.pt: stock.csv bytes plus the catalog id list.
+
+    Image sizes are not included, so another view of the same SKU does not
+    invalidate a head. A changed row or id does. None if the id list is unsafe.
+    """
+    cleaned = head_class_set(sku_ids)
+    if cleaned is None:
+        return None
+    digest = hashlib.sha256()
+    csv_path = os.path.join(stock_dir, csv_name)
+    try:
+        with open(csv_path, "rb") as handle:
+            digest.update(handle.read())
+    except OSError:
+        digest.update(b"missing-csv")
+    digest.update(b"\n")
+    digest.update("\n".join(sorted(cleaned)).encode("utf-8"))
+    return digest.hexdigest()
+
+
+
+
 class ProductGallery:
     """Several views per SKU, embedded once, queried by cosine similarity."""
 
@@ -162,44 +196,85 @@ class ProductGallery:
         # 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_reason = "not_loaded"
         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.
+        """Load the store's fine-tuned head when its labels are this catalog's.
 
         weights_only: the file is tensors and strings, not executable code.
-        A payload that is not that head leaves the raw gallery in place.
+        A missing, malformed, or foreign head leaves the raw gallery in place
+        and records head_reason. Class ids must be unique, nonempty, and a
+        subset of this gallery: a head may cover only the SKUs it was trained
+        on. A catalog_digest, when present, must match stock.csv plus that
+        subset. A head written before the pin loads with head_reason
+        legacy_unpinned; it still cannot name a SKU outside this catalog.
         """
         import torch
 
+        self._head = None
         if not os.path.exists(self._head_path):
+            self.head_reason = "missing"
             return False
         try:
             payload = torch.load(self._head_path, map_location="cpu", weights_only=True)
-            if not isinstance(payload, dict):
+            reason = self._accept_head(payload, torch)
+        except Exception:
+            self._head = None
+            self.head_reason = "malformed"
             return False
+        if not reason.startswith("ok"):
+            self._head = None
+            self.head_reason = reason
+            return False
+        self.head_reason = "legacy_unpinned" if reason == "ok:legacy_unpinned" else "ok"
+        return True
+
+    def _accept_head(self, payload: Any, torch: Any) -> str:
+        if not isinstance(payload, dict):
+            return "malformed"
         projection_state = payload.get("projection")
         prototypes = payload.get("prototypes")
         class_ids = payload.get("class_ids")
         if not isinstance(projection_state, dict) or not torch.is_tensor(prototypes):
-                return False
+            return "malformed"
         if not isinstance(class_ids, (list, tuple)) or prototypes.ndim != 2:
-                return False
+            return "malformed"
-            if prototypes.shape[0] != len(class_ids):
+        if prototypes.shape[0] != len(class_ids) or prototypes.shape[1] != 128:
-                return False
+            return "malformed"
+        labels = []
+        for sku_id in class_ids:
+            if not isinstance(sku_id, str) or not sku_id.strip():
+                return "malformed: identity head class_ids must be nonempty strings"
+            labels.append(sku_id.strip())
+        if len(set(labels)) != len(labels):
+            return "malformed: identity head class_ids are not unique"
+        expected = head_class_set(self.sku_ids)
+        if expected is None or not set(labels) <= set(expected):
+            return "catalog_mismatch: head class_ids are not a subset of this catalog; retrain finetune_identity.py --stock"
+        recorded = payload.get("catalog_digest", None)
+        if recorded is None:
+            reason = "ok:legacy_unpinned"
+        else:
+            pinned = catalog_digest(self.stock_dir, labels)
+            if pinned is None or recorded != pinned:
+                return "catalog_mismatch: head catalog_digest does not match stock.csv and its class_ids; retrain finetune_identity.py --stock"
+            reason = "ok"
+        if payload.get("base_model") != DINOV2_MODEL:
+            return f"backbone_mismatch: head base_model is not {DINOV2_MODEL}"
+        if type(payload.get("image_size")) is not int or payload.get("image_size") != self.image_size:
+            return f"image_size_mismatch: head image_size is not {self.image_size}"
         projection = torch.nn.Linear(384, 128)
         projection.load_state_dict(projection_state)
         projection.eval()
         self._head = {
             "projection": projection,
             "prototypes": torch.nn.functional.normalize(prototypes.float(), dim=1),
-                "class_ids": [str(sku_id) for sku_id in class_ids],
+            "class_ids": labels,
         }
-        except Exception:
+        return reason
-            self._head = None
+
-            return False
-        return True
 
     def _load_model(self) -> Any:
         if self._model is not None:

recognization/code/test_gallery.py

@@ -22,10 +22,13 @@ if current_dir not in sys.path:
     sys.path.insert(0, current_dir)
 
 from gallery import (  # noqa: E402
+    DINOV2_MODEL,
     DINOV2_SHA256,
     GalleryEntry,
     ProductGallery,
+    catalog_digest,
     dinov2_weights_path,
+    head_class_set,
     load_dinov2,
 )
 
@@ -144,5 +147,83 @@ class OfflineWeightsTests(unittest.TestCase):
             if previous is not None:
                 os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
 
+
+def _head_payload(class_ids, torch):
+    projection = torch.nn.Linear(384, 128)
+    prototypes = torch.nn.functional.normalize(torch.randn(len(class_ids), 128), dim=1)
+    return {
+        "class_ids": list(class_ids),
+        "image_size": 518,
+        "base_model": DINOV2_MODEL,
+        "projection": {key: value.detach().cpu() for key, value in projection.state_dict().items()},
+        "prototypes": prototypes.detach().cpu(),
+    }
+
+
+class HeadPinTests(unittest.TestCase):
+    def test_a_stale_head_cannot_label_this_catalog(self):
+        import tempfile
+
+        import torch
+
+        with tempfile.TemporaryDirectory() as directory:
+            with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
+                handle.write("SKU_ID,product_name\nsku_a,A\nsku_b,B\nsku_c,C\n")
+            gallery = ProductGallery(directory, ["sku_a", "sku_b", "sku_c"])
+            gallery._available = True
+            path = gallery._head_path
+            subset = ["sku_c", "sku_a"]
+            digest = catalog_digest(directory, subset)
+
+            foreign = _head_payload(["other_store", "sku_a"], torch)
+            torch.save(foreign, path)
+            self.assertFalse(gallery.load_head())
+            self.assertIsNone(gallery._head)
+            self.assertIn("catalog_mismatch", gallery.head_reason)
+            self.assertTrue(gallery.available())
+
+            legacy = _head_payload(subset, torch)
+            torch.save(legacy, path)
+            self.assertTrue(gallery.load_head())
+            self.assertEqual(gallery.head_reason, "legacy_unpinned")
+            self.assertEqual(gallery._head["class_ids"], subset)
+            self.assertNotIn("other_store", gallery._head["class_ids"])
+
+            stale = _head_payload(subset, torch)
+            stale["catalog_digest"] = "0" * 64
+            torch.save(stale, path)
+            self.assertFalse(gallery.load_head())
+            self.assertIsNone(gallery._head)
+            self.assertIn("catalog_mismatch", gallery.head_reason)
+            self.assertTrue(gallery.available())
+
+            wrong_backbone = _head_payload(subset, torch)
+            wrong_backbone["catalog_digest"] = digest
+            wrong_backbone["base_model"] = "vit_base_patch14_dinov2.lvd142m"
+            torch.save(wrong_backbone, path)
+            self.assertFalse(gallery.load_head())
+            self.assertIn("backbone_mismatch", gallery.head_reason)
+
+            wrong_size = _head_payload(subset, torch)
+            wrong_size["catalog_digest"] = digest
+            wrong_size["image_size"] = 224
+            torch.save(wrong_size, path)
+            self.assertFalse(gallery.load_head())
+            self.assertIn("image_size_mismatch", gallery.head_reason)
+
+            torch.save(["not", "a", "head"], path)
+            self.assertFalse(gallery.load_head())
+            self.assertTrue(gallery.head_reason.startswith("malformed"))
+            self.assertTrue(gallery.available())
+            self.assertIsNone(gallery._head)
+
+            current = _head_payload(subset, torch)
+            current["catalog_digest"] = digest
+            torch.save(current, path)
+            self.assertTrue(gallery.load_head())
+            self.assertEqual(gallery.head_reason, "ok")
+            self.assertEqual(gallery._head["class_ids"], subset)
+            self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))
+
 if __name__ == "__main__":
     unittest.main()