fix: constrain trained head labels to enrolled catalog

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,45 +196,86 @@ 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):
-                return False
-            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
-            if not isinstance(class_ids, (list, tuple)) or prototypes.ndim != 2:
-                return False
-            if prototypes.shape[0] != len(class_ids):
-                return False
-            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],
-            }
+            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 "malformed"
+        if not isinstance(class_ids, (list, tuple)) or prototypes.ndim != 2:
+            return "malformed"
+        if prototypes.shape[0] != len(class_ids) or prototypes.shape[1] != 128:
+            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": labels,
+        }
+        return reason
+
+
     def _load_model(self) -> Any:
         if self._model is not None:
             return self._model