From ac315eb48e8fe7a5f10dfd7d89fa2e080efe0f42 Mon Sep 17 00:00:00 2001
From: kiyo9w <ngokapikapi@gmail.com>
Date: Mon, 28 Sep 2026 14:41:57 +0900
Subject: [PATCH] fix: refuse unavailable gallery and pin head trainer

---
 recognization/code/app.py               | 10 +++++--
 recognization/code/finetune_identity.py |  8 ++++--
 recognization/code/test_engine.py       | 27 ++++++++++++++++++
 recognization/code/test_gallery.py      | 38 +++++++++++++++++++++++++
 4 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/recognization/code/app.py b/recognization/code/app.py
index abb1e8c..c939471 100644
--- a/recognization/code/app.py
+++ b/recognization/code/app.py
@@ -293,8 +293,11 @@ def create_app(
 
     scans = scans if scans is not None else build_legacy_scans(engine)
     # A live session can be recorded so a real scan becomes data to label later.
-    recorder = FrameRecorder(settings.record_frames, every=settings.record_every) \
-        if settings.record_frames else None
+    recorder = FrameRecorder(
+        settings.record_frames,
+        every=settings.record_every,
+        identity_backend=engine.settings.identity_backend,
+    ) if settings.record_frames else None
     samples = {item["id"]: item for item in get_available_shelf_images("add")}
     # path -> (mtime_ns, size, width, height). Dropped when the file changes.
     sample_dimensions: Dict[str, Tuple[int, int, int, int]] = {}
@@ -1138,7 +1141,8 @@ def main() -> None:
     engine = build_default_engine(engine_settings)
     if (engine_settings.identity_backend or "").lower() == "gallery":
         if engine.identity_gallery() is None:
-            raise RuntimeError(
+            engine.shutdown()
+            raise SystemExit(
                 "gallery identity backend requires local DINOv2 weights; refusing to start"
             )
     warm_ms = engine.warm_up()
diff --git a/recognization/code/finetune_identity.py b/recognization/code/finetune_identity.py
index fc29d73..d31741a 100644
--- a/recognization/code/finetune_identity.py
+++ b/recognization/code/finetune_identity.py
@@ -132,12 +132,14 @@ def enrollment_report_text(report: dict) -> str:
 
 
 def build_model(image_size: int = 518):
-    import timm
     import torch
 
+    from gallery import load_dinov2
+
     device = "mps" if torch.backends.mps.is_available() else "cpu"
-    model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
-                              num_classes=0, img_size=image_size, dynamic_img_size=True)
+    # Same hash-pinned local file as inference. pretrained=True would fetch a
+    # different copy even with YOLO_OFFLINE=1.
+    model = load_dinov2(image_size)
     return model.eval().to(device), device
 
 
diff --git a/recognization/code/test_engine.py b/recognization/code/test_engine.py
index d32ac67..7eb686f 100644
--- a/recognization/code/test_engine.py
+++ b/recognization/code/test_engine.py
@@ -274,6 +274,33 @@ class CatalogTests(unittest.TestCase):
         self.assertEqual(product["number_out"], 3)
         self.assertIsNone(engine.update_stock("nope", -1))
 
+
+class GalleryLoadTests(unittest.TestCase):
+    def test_a_failed_gallery_load_returns_none_and_is_not_retried(self):
+        engine, _, _, _ = make_engine(identity_backend="gallery", stock_dir="does-not-exist")
+
+        class BrokenGallery:
+            def __init__(self, *args, **kwargs):
+                pass
+
+            def load(self):
+                raise FileNotFoundError("DINOv2 weights not found")
+
+        import gallery as gallery_module
+
+        original = gallery_module.ProductGallery
+        gallery_module.ProductGallery = BrokenGallery
+        try:
+            self.assertIsNone(engine.identity_gallery())
+            self.assertIsNone(engine.identity_gallery())
+        finally:
+            gallery_module.ProductGallery = original
+            engine.shutdown()
+        self.assertIs(engine._gallery, False)
+
+
+
+class CatalogImageTests(unittest.TestCase):
     def test_catalog_image_path_is_resolved_through_inventory(self):
         engine, _, _, _ = make_engine()
         self.assertIsNone(engine.catalog_image_path("../../etc/passwd"))
diff --git a/recognization/code/test_gallery.py b/recognization/code/test_gallery.py
index 477d4f8..3c4ddba 100644
--- a/recognization/code/test_gallery.py
+++ b/recognization/code/test_gallery.py
@@ -225,5 +225,43 @@ class HeadPinTests(unittest.TestCase):
             self.assertEqual(gallery._head["class_ids"], subset)
             self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))
 
+
+class TrainerPinTests(unittest.TestCase):
+    def test_head_training_loads_the_pinned_file_and_keeps_preprocess(self):
+        import finetune_identity
+        import gallery as gallery_module
+
+        seen = []
+
+        class FakeBackbone:
+            def eval(self):
+                return self
+
+            def to(self, device):
+                return self
+
+        sentinel = FakeBackbone()
+        def fake_load(image_size):
+            seen.append(image_size)
+            return sentinel
+
+
+        original = finetune_identity.load_dinov2
+        finetune_identity.load_dinov2 = fake_load
+        try:
+            model, device = finetune_identity.build_model(224)
+        finally:
+            finetune_identity.load_dinov2 = original
+        self.assertIs(model, sentinel)
+        self.assertIn(device, ("cpu", "mps"))
+        self.assertEqual(seen, [224])
+
+        image = np.zeros((10, 20, 3), dtype=np.uint8)
+        self.assertEqual(finetune_identity.prepare(image, 32).shape, gallery_module._prepare(image, 32).shape)
+        self.assertTrue(np.allclose(
+            finetune_identity.prepare(image, 32),
+            gallery_module._prepare(image, 32),
+        ))
+
 if __name__ == "__main__":
     unittest.main()