diff --git a/recognization/code/finetune_identity.py b/recognization/code/finetune_identity.py index c147a1d..fc29d73 100644 --- a/recognization/code/finetune_identity.py +++ b/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())}, diff --git a/recognization/code/gallery.py b/recognization/code/gallery.py index deeb986..1272e3c 100644 --- a/recognization/code/gallery.py +++ b/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 diff --git a/recognization/code/test_gallery.py b/recognization/code/test_gallery.py index 815365e..477d4f8 100644 --- a/recognization/code/test_gallery.py +++ b/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()