diff --git a/recognization/code/app.py b/recognization/code/app.py index abb1e8c..bcbd1ff 100644 --- a/recognization/code/app.py +++ b/recognization/code/app.py @@ -21,7 +21,6 @@ import argparse import asyncio import base64 from contextlib import asynccontextmanager -from dataclasses import dataclass import json import logging import os @@ -30,7 +29,8 @@ import sys import threading import time import uuid -from typing import Any, Dict, List, Optional, Tuple +from dataclasses import dataclass +from typing import Any, Dict, List, Optional # Keep the demo off the network before anything imports ultralytics. os.environ.setdefault("YOLO_OFFLINE", "1") @@ -296,9 +296,6 @@ 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 @@ -374,37 +371,15 @@ def create_app( raise _error("NOT_FOUND", f"unknown sku '{sku_id}'", 404) return {"status": "success", "product": product} - 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 = [] + import cv2 + for item in samples.values(): entry = {"id": item["id"], "label": item["label"], "url": f"/v1/samples/{item['id']}"} - dimensions = _sample_dimensions(item["path"]) - if dimensions is not None: - entry["width"], entry["height"] = dimensions + image = cv2.imread(item["path"]) + if image is not None: + entry["width"], entry["height"] = int(image.shape[1]), int(image.shape[0]) listing.append(entry) return listing @@ -718,6 +693,7 @@ 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", @@ -741,7 +717,7 @@ def create_app( async def processing_loop() -> None: while not session.closed: - pending = await session.take_pending_async() + pending = await loop.run_in_executor(None, session.take_pending, 0.25) if pending is None: continue frame_id, payload = pending diff --git a/recognization/code/engine.py b/recognization/code/engine.py index 6a84c92..f63e714 100644 --- a/recognization/code/engine.py +++ b/recognization/code/engine.py @@ -273,10 +273,6 @@ 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 @@ -481,8 +477,6 @@ 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 @@ -612,53 +606,24 @@ class RecognitionEngine: # ------------------------------------------------------------------ catalog def catalog_version(self) -> str: - """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. - """ + digest = hashlib.sha256() csv_path = os.path.join(self.settings.stock_dir, self.settings.stock_csv) try: - stat = os.stat(csv_path) - csv_key = (stat.st_mtime_ns, stat.st_size) + with open(csv_path, "rb") as handle: + digest.update(handle.read()) except OSError: - csv_key = None + digest.update(b"missing-csv") try: - entries = tuple(sorted( + entries = 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 = () - 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 + entries = [] for name, size in entries: digest.update(f"{name}:{size}".encode("utf-8")) - 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 + return digest.hexdigest()[:12] def model_version(self) -> str: try: @@ -724,8 +689,6 @@ 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: diff --git a/recognization/code/finetune_identity.py b/recognization/code/finetune_identity.py index fc29d73..c147a1d 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 DINOV2_MODEL, ProductGallery, catalog_digest, head_class_set + from gallery import ProductGallery stock = args.stock report = enrollment_report(stock) @@ -267,15 +267,10 @@ 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": DINOV2_MODEL, + "base_model": "vit_small_patch14_dinov2.lvd142m", "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 1272e3c..deeb986 100644 --- a/recognization/code/gallery.py +++ b/recognization/code/gallery.py @@ -25,7 +25,6 @@ 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 @@ -132,39 +131,6 @@ 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.""" @@ -196,86 +162,45 @@ 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 when its labels are this catalog's. + """Load the store's fine-tuned head, if it has one. weights_only: the file is tensors and strings, not executable code. - 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. + A payload that is not that head leaves the raw gallery in place. """ 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) - reason = self._accept_head(payload, torch) + 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], + } 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/live.py b/recognization/code/live.py index 3850448..385f379 100644 --- a/recognization/code/live.py +++ b/recognization/code/live.py @@ -13,7 +13,6 @@ lock so classification callbacks can land concurrently. from __future__ import annotations -import asyncio import collections import os import statistics @@ -170,9 +169,6 @@ 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 @@ -217,77 +213,23 @@ 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: - 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() + pending = self._pending + self._pending = None + self._frame_ready.clear() + return pending def close(self) -> None: with self._lock: self._closed = True self._pending = None - self._frame_ready.set() - self._wake_waiter_locked() + self._frame_ready.set() @property def closed(self) -> bool: @@ -351,8 +293,6 @@ 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_gallery.py b/recognization/code/test_gallery.py index 477d4f8..815365e 100644 --- a/recognization/code/test_gallery.py +++ b/recognization/code/test_gallery.py @@ -22,13 +22,10 @@ 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, ) @@ -147,83 +144,5 @@ 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() diff --git a/recognization/code/test_live.py b/recognization/code/test_live.py index 3fbdc7d..1e9ed14 100644 --- a/recognization/code/test_live.py +++ b/recognization/code/test_live.py @@ -11,7 +11,6 @@ newest frame. from __future__ import annotations import os -import asyncio import sys import threading import time @@ -392,39 +391,6 @@ 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()