From 1064e356de94aca24ea83a02c6d9e4673f3eacce Mon Sep 17 00:00:00 2001 From: kiyo9w Date: Mon, 28 Sep 2026 14:06:34 +0900 Subject: [PATCH 1/3] perf: reuse sample dimensions and catalog digest --- recognization/code/app.py | 42 ++++++++++++++++++++++------ recognization/code/engine.py | 53 ++++++++++++++++++++++++++++++------ 2 files changed, 78 insertions(+), 17 deletions(-) diff --git a/recognization/code/app.py b/recognization/code/app.py index bcbd1ff..abb1e8c 100644 --- a/recognization/code/app.py +++ b/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 diff --git a/recognization/code/engine.py b/recognization/code/engine.py index f63e714..6a84c92 100644 --- a/recognization/code/engine.py +++ b/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: From b937be8e30121124568c7227f7a05515ebc09105 Mon Sep 17 00:00:00 2001 From: kiyo9w Date: Mon, 28 Sep 2026 14:06:57 +0900 Subject: [PATCH 2/3] perf: await pending frames without idle worker --- recognization/code/live.py | 70 ++++++++++++++++++++++++++++++--- recognization/code/test_live.py | 34 ++++++++++++++++ 2 files changed, 99 insertions(+), 5 deletions(-) diff --git a/recognization/code/live.py b/recognization/code/live.py index 385f379..3850448 100644 --- a/recognization/code/live.py +++ b/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]: diff --git a/recognization/code/test_live.py b/recognization/code/test_live.py index 1e9ed14..3fbdc7d 100644 --- a/recognization/code/test_live.py +++ b/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() From 27983224177ef2e39c4229b08242e420743b54da Mon Sep 17 00:00:00 2001 From: kiyo9w Date: Mon, 28 Sep 2026 14:07:02 +0900 Subject: [PATCH 3/3] fix: constrain trained head labels to enrolled catalog --- recognization/code/finetune_identity.py | 9 +- recognization/code/gallery.py | 117 +++++++++++++++++++----- recognization/code/test_gallery.py | 81 ++++++++++++++++ 3 files changed, 184 insertions(+), 23 deletions(-) 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()