fix: constrain trained head labels to enrolled catalog

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()