fix: constrain trained head labels to enrolled catalog

recognization/code/app.py

@@ -21,6 +21,7 @@ import argparse
 import asyncio
 import base64
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 import json
 import logging
 import os
@@ -29,8 +30,7 @@ import sys
 import threading
 import time
 import uuid
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 # Keep the demo off the network before anything imports ultralytics.
 os.environ.setdefault("YOLO_OFFLINE", "1")
@@ -296,6 +296,9 @@ def create_app(
     recorder = FrameRecorder(settings.record_frames, every=settings.record_every) \
         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]] = {}
+    sample_dimensions_lock = threading.Lock()
     sessions: Dict[str, LiveSession] = {}
     sessions_lock = asyncio.Lock()
     app.state.engine = engine
@@ -371,15 +374,37 @@ def create_app(
             raise _error("NOT_FOUND", f"unknown sku '{sku_id}'", 404)
         return {"status": "success", "product": product}
 
-    def _sample_listing() -> List[Dict[str, Any]]:
-        listing = []
+    def _sample_dimensions(path: str) -> Optional[Tuple[int, int]]:
+        try:
+            stat = os.stat(path)
+        except OSError:
+            with sample_dimensions_lock:
+                sample_dimensions.pop(path, None)
+            return None
+        key = (stat.st_mtime_ns, stat.st_size)
+        with sample_dimensions_lock:
+            cached = sample_dimensions.get(path)
+            if cached is not None and cached[:2] == key:
+                return cached[2], cached[3]
         import cv2
 
+        image = cv2.imread(path)
+        if image is None:
+            with sample_dimensions_lock:
+                sample_dimensions.pop(path, None)
+            return None
+        width, height = int(image.shape[1]), int(image.shape[0])
+        with sample_dimensions_lock:
+            sample_dimensions[path] = (key[0], key[1], width, height)
+        return width, height
+
+    def _sample_listing() -> List[Dict[str, Any]]:
+        listing = []
         for item in samples.values():
             entry = {"id": item["id"], "label": item["label"], "url": f"/v1/samples/{item['id']}"}
-            image = cv2.imread(item["path"])
-            if image is not None:
-                entry["width"], entry["height"] = int(image.shape[1]), int(image.shape[0])
+            dimensions = _sample_dimensions(item["path"])
+            if dimensions is not None:
+                entry["width"], entry["height"] = dimensions
             listing.append(entry)
         return listing
 
@@ -693,7 +718,6 @@ def create_app(
             async def send(payload: Dict[str, Any]) -> None:
                 async with send_lock:
                     await websocket.send_text(json.dumps(payload))
-
             await send(
                 {
                     "type": "hello",
@@ -717,7 +741,7 @@ def create_app(
 
         async def processing_loop() -> None:
             while not session.closed:
-                pending = await loop.run_in_executor(None, session.take_pending, 0.25)
+                pending = await session.take_pending_async()
                 if pending is None:
                     continue
                 frame_id, payload = pending

recognization/code/engine.py

@@ -273,6 +273,10 @@ class RecognitionEngine:
             self.settings.image_store_max_bytes,
         )
         self._catalog_cache: Optional[Tuple[str, List[Dict[str, Any]]]] = None
+        # (version, csv mtime_ns, csv size, sorted image name:size). Version is
+        # the content digest; the rest is the cheap fingerprint that invalidates it.
+        self._catalog_version_cache: Optional[Tuple[str, int, int, Tuple[Tuple[str, int], ...]]] = None
+        self._catalog_version_lock = threading.Lock()
         self._started_at = time.time()
 
     # ------------------------------------------------------------------ lifecycle
@@ -477,6 +481,8 @@ class RecognitionEngine:
                     len(gallery.sku_ids),
                     gallery.device,
                 )
+                if gallery.head_reason != "ok":
+                    LOGGER.warning("identity head: %s (%s)", gallery.head_reason, gallery._head_path)
             else:
                 LOGGER.warning("identity gallery unavailable; gallery requests stay unrecognized")
             return gallery if loaded else None
@@ -606,24 +612,53 @@ class RecognitionEngine:
 
     # ------------------------------------------------------------------ catalog
     def catalog_version(self) -> str:
-        digest = hashlib.sha256()
+        """12-hex digest of stock.csv bytes plus template image names and sizes.
+
+        Unchanged files reuse the digest. A stock update that has not been written
+        back to the CSV still changes the digest, because the in-memory fingerprint
+        is dropped. Same-size content rewritten without a new mtime is not detected
+        until the next stock update or a real stat change.
+        """
         csv_path = os.path.join(self.settings.stock_dir, self.settings.stock_csv)
         try:
-            with open(csv_path, "rb") as handle:
-                digest.update(handle.read())
+            stat = os.stat(csv_path)
+            csv_key = (stat.st_mtime_ns, stat.st_size)
         except OSError:
-            digest.update(b"missing-csv")
+            csv_key = None
         try:
-            entries = sorted(
+            entries = tuple(sorted(
                 (name, os.path.getsize(os.path.join(self.settings.stock_dir, name)))
                 for name in os.listdir(self.settings.stock_dir)
                 if name.lower().endswith((".jpeg", ".jpg", ".png"))
-            )
+            ))
         except OSError:
-            entries = []
+            entries = ()
+        fingerprint = (
+            csv_key[0] if csv_key else -1,
+            csv_key[1] if csv_key else -1,
+            entries,
+        )
+        with self._catalog_version_lock:
+            cached = self._catalog_version_cache
+            if cached is not None and csv_key is not None and cached[1:] == fingerprint:
+                return cached[0]
+        digest = hashlib.sha256()
+        if csv_key is None:
+            digest.update(b"missing-csv")
+        else:
+            try:
+                with open(csv_path, "rb") as handle:
+                    digest.update(handle.read())
+            except OSError:
+                digest.update(b"missing-csv")
+                csv_key = None
         for name, size in entries:
             digest.update(f"{name}:{size}".encode("utf-8"))
-        return digest.hexdigest()[:12]
+        version = digest.hexdigest()[:12]
+        if csv_key is not None:
+            with self._catalog_version_lock:
+                self._catalog_version_cache = (version, csv_key[0], csv_key[1], entries)
+        return version
 
     def model_version(self) -> str:
         try:
@@ -689,6 +724,8 @@ class RecognitionEngine:
             product = self.inventory.get_product(sku_id) or {}
             payload = self._product_payload(sku_id, product)
         self._catalog_cache = None
+        with self._catalog_version_lock:
+            self._catalog_version_cache = None
         return payload
 
     def sku_count(self) -> int:

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,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

recognization/code/live.py

@@ -13,6 +13,7 @@ lock so classification callbacks can land concurrently.
 
 from __future__ import annotations
 
+import asyncio
 import collections
 import os
 import statistics
@@ -169,6 +170,9 @@ class LiveSession:
 
         self._pending: Optional[Tuple[int, bytes]] = None
         self._frame_ready = threading.Event()
+        # One waiter, created on the caller's loop. None until the first async take.
+        self._waiter: Optional[asyncio.Event] = None
+        self._waiter_loop: Optional[asyncio.AbstractEventLoop] = None
         self._closed = False
 
         self._seq = 0
@@ -213,23 +217,77 @@ class LiveSession:
                 self._frames_dropped += 1
             self._pending = (int(frame_id), payload)
             self._frame_ready.set()
+            self._wake_waiter_locked()
             return True
 
+    def _pop_pending_locked(self) -> Optional[Tuple[int, bytes]]:
+        pending = self._pending
+        self._pending = None
+        self._frame_ready.clear()
+        return pending
+
+    def _wake_waiter_locked(self) -> None:
+        """Set the async event from this thread. Caller holds self._lock."""
+        waiter = self._waiter
+        loop = self._waiter_loop
+        if waiter is None or loop is None:
+            return
+        try:
+            running = asyncio.get_running_loop()
+        except RuntimeError:
+            running = None
+        if running is loop:
+            waiter.set()
+            return
+        try:
+            loop.call_soon_threadsafe(waiter.set)
+        except RuntimeError:
+            # Loop already closed. The sync event still wakes take_pending.
+            pass
+
     def take_pending(self, timeout: float = 0.2) -> Optional[Tuple[int, bytes]]:
         """Blocking pop for the processing loop; None on timeout/close."""
         if not self._frame_ready.wait(timeout):
             return None
         with self._lock:
-            pending = self._pending
-            self._pending = None
-            self._frame_ready.clear()
-        return pending
+            return self._pop_pending_locked()
+
+    async def take_pending_async(self) -> Optional[Tuple[int, bytes]]:
+        """Awaitable pop. None means the session closed with nothing left.
+
+        Bind the waiter to the running loop on first use. A frame that arrives
+        before that bind is still visible: the event starts set when one is queued.
+        """
+        loop = asyncio.get_running_loop()
+        with self._lock:
+            if self._waiter is not None and self._waiter_loop is not loop:
+                raise RuntimeError("take_pending_async is bound to another event loop")
+            if self._waiter is None:
+                self._waiter = asyncio.Event()
+                self._waiter_loop = loop
+                if self._pending is not None or self._closed:
+                    self._waiter.set()
+            waiter = self._waiter
+
+        while True:
+            await waiter.wait()
+            with self._lock:
+                if self._pending is not None:
+                    pending = self._pop_pending_locked()
+                    if self._waiter is waiter and self._pending is None and not self._closed:
+                        self._waiter.clear()
+                    return pending
+                if self._closed:
+                    return None
+                if self._waiter is waiter:
+                    self._waiter.clear()
 
     def close(self) -> None:
         with self._lock:
             self._closed = True
             self._pending = None
-        self._frame_ready.set()
+            self._frame_ready.set()
+            self._wake_waiter_locked()
 
     @property
     def closed(self) -> bool:
@@ -293,6 +351,8 @@ class LiveSession:
             self._latest_frame_id = None
             self._frame_origin = [0.0, 0.0]
             self._shift_map = None
+            # ponytail: do not touch _waiter here. asyncio.Event is loop-affine;
+            # take_pending_async clears a stale wake on its own loop.
 
     # ------------------------------------------------------------------ frame path
     def process_frame(self, frame_id: int, payload: bytes) -> Dict[str, Any]:

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

recognization/code/test_live.py

@@ -11,6 +11,7 @@ newest frame.
 from __future__ import annotations
 
 import os
+import asyncio
 import sys
 import threading
 import time
@@ -391,6 +392,39 @@ class ClassificationBudgetTests(unittest.TestCase):
         session.close()
         self.assertTrue(session.closed)
         self.assertIsNone(session.take_pending(timeout=0.05))
+
+    def test_async_take_keeps_latest_and_wakes_on_close(self):
+        """The server waiter must not miss a frame that beat registration, and close must wake it."""
+        session, engine, _, _ = make_session()
+
+        async def scenario():
+            session.submit_frame(1, frame_bytes())
+            session.submit_frame(2, frame_bytes())
+            first = await asyncio.wait_for(session.take_pending_async(), timeout=1.0)
+            self.assertEqual(first[0], 2)
+
+            waiting = asyncio.create_task(session.take_pending_async())
+            await asyncio.sleep(0)
+            self.assertFalse(waiting.done(), "an empty queue must not complete the waiter")
+            session.reset()
+            session.submit_frame(3, frame_bytes())
+            self.assertEqual((await asyncio.wait_for(waiting, timeout=1.0))[0], 3)
+
+            cancelled = asyncio.create_task(session.take_pending_async())
+            await asyncio.sleep(0)
+            cancelled.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await cancelled
+            session.submit_frame(4, frame_bytes())
+            self.assertEqual((await asyncio.wait_for(session.take_pending_async(), timeout=1.0))[0], 4)
+
+            closing = asyncio.create_task(session.take_pending_async())
+            await asyncio.sleep(0)
+            session.close()
+            self.assertIsNone(await asyncio.wait_for(closing, timeout=1.0))
+            self.assertFalse(session.submit_frame(5, frame_bytes()))
+
+        asyncio.run(scenario())
         engine.shutdown()