""" Recognition engine for the Dorayaki shelf-recognition server. One process-wide, thread-safe engine. Every call is stateless from the caller's point of view: the active image is never stored here, so any number of clients can call concurrently. Owns: - the YOLO product detector (single instance, every predict under a lock), - the feature-matching classifier (thread-local OpenCV extractors, a classify thread pool), - the stock catalog (InventoryManager, mutations under a lock), - empty-slot geometry (resolution-normalized), - a small LRU store of recent still images for later classify/inspect calls. Heavy imports (torch, ultralytics, cv2) are pulled in lazily by build_default_engine so tests can inject fakes without loading weights. """ from __future__ import annotations import base64 import collections import hashlib import json import os import threading import time from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace from typing import Any, Dict, List, Optional, Sequence, Tuple # Keep the demo box off the network before ultralytics is imported anywhere. os.environ.setdefault("YOLO_OFFLINE", "1") os.environ.setdefault("YOLO_AUTOINSTALL", "false") COORDINATE_CONVENTION = "pixel-xyxy-origin-top-left" REFERENCE_WIDTH = 2048.0 MAX_UPSCALE = 2.0 DEFAULT_IMAGE_STORE_MAX = 16 DEFAULT_IMAGE_STORE_TTL_S = 15 * 60 DEFAULT_IMAGE_STORE_MAX_BYTES = 256 * 1024 * 1024 STATE_DETECTED = "detected" STATE_IDENTIFYING = "identifying" STATE_RECOGNIZED = "recognized" STATE_NEEDS_CONFIRMATION = "needs_confirmation" STATE_UNRECOGNIZED = "unrecognized" CLASSIFY_METHODS = ("surf", "orb", "brisk", "hybrid", "hybrid_fast") class UndecodableImage(ValueError): """Raised when an uploaded payload is not a readable image.""" class UnknownMethod(ValueError): """Raised for a classification method outside CLASSIFY_METHODS.""" class ImageExpired(KeyError): """Raised when an image_id is no longer in the store.""" @dataclass class EngineSettings: weights_path: str = "runs/detect/train/weights/best.pt" stock_dir: str = "stock" stock_csv: str = "stock.csv" # Measured on this project (Apple M5, bench_engine.py): CPU detects in ~27 ms # p50 while MPS takes ~112 ms with identical boxes, so CPU is the default. # Override with --device / DORAYAKI_DEVICE (cpu | mps | cuda | cuda:0). device: Optional[str] = "cpu" # Live-mode defaults, measured with stream_client.py --pace result: # hybrid classifies at ~44 ms/crop p50 and is the most accurate method # (21/21 SKUs on image.jpg). hybrid_fast (ORB shortlist) is 3.4x faster but # agreed with hybrid on only 47% of accepted crops, so it is NOT the default # and stays available for experiments only (bench_engine.py --compare). live_inflight: int = 12 imgsz: int = 640 classify_workers: int = 0 # 0 -> max(2, cpu_count - 2) live_method: str = "hybrid" live_conf: float = 0.1 detect_conf: float = 0.25 detect_iou: float = 0.45 image_store_max: int = DEFAULT_IMAGE_STORE_MAX image_store_ttl_s: float = DEFAULT_IMAGE_STORE_TTL_S image_store_max_bytes: int = DEFAULT_IMAGE_STORE_MAX_BYTES @dataclass class StoredImage: image_id: str image: Any # numpy BGR array width: int height: int detections: List[Dict[str, Any]] = field(default_factory=list) created_at: float = 0.0 nbytes: int = 0 def identity_state_from_level(level: Optional[str]) -> str: """Map a decide_identity() level to the wire state vocabulary.""" if level == "accepted": return STATE_RECOGNIZED if level in ("ambiguous", "unknown"): return STATE_NEEDS_CONFIRMATION return STATE_UNRECOGNIZED def identity_payload(decision: Optional[Dict[str, Any]], *, attempts: int = 1) -> Dict[str, Any]: """One wire-shaped identity object from a decide_identity() result.""" if not decision: return { "state": STATE_DETECTED, "sku_id": None, "identity_level": None, "confidence": None, "score": None, "candidates": [], "attempts": 0, "method": None, } level = decision.get("identity_level") state = identity_state_from_level(level) candidates = [ {"sku_id": row.get("sku_id"), "score": row.get("score")} for row in (decision.get("ranked_candidates") or [])[:3] ] return { "state": state, "sku_id": decision.get("best_sku") if state == STATE_RECOGNIZED else None, "identity_level": level, "confidence": decision.get("confidence"), "score": decision.get("score"), "candidates": candidates, "attempts": attempts, "method": decision.get("method"), } class ImageStore: """Bounded LRU of decoded still images keyed by a random image_id.""" def __init__(self, max_items: int = DEFAULT_IMAGE_STORE_MAX, ttl_s: float = DEFAULT_IMAGE_STORE_TTL_S, max_bytes: int = DEFAULT_IMAGE_STORE_MAX_BYTES): self.max_items = max(1, int(max_items)) self.ttl_s = float(ttl_s) self.max_bytes = max(1, int(max_bytes)) self._items: "collections.OrderedDict[str, StoredImage]" = collections.OrderedDict() self._bytes = 0 self._lock = threading.Lock() def put(self, image: Any, detections: Optional[List[Dict[str, Any]]] = None) -> StoredImage: height, width = image.shape[:2] stored = StoredImage( image_id="img_" + os.urandom(8).hex(), image=image, width=int(width), height=int(height), detections=list(detections or []), created_at=time.time(), nbytes=int(image.nbytes), ) with self._lock: self._evict_locked() self._items[stored.image_id] = stored self._bytes += stored.nbytes while len(self._items) > self.max_items or self._bytes > self.max_bytes: if len(self._items) <= 1: break _, dropped = self._items.popitem(last=False) self._bytes -= dropped.nbytes return stored def get(self, image_id: str) -> StoredImage: with self._lock: stored = self._items.get(image_id) if stored is None or time.time() - stored.created_at > self.ttl_s: if stored is not None: self._items.pop(image_id, None) self._bytes -= stored.nbytes raise ImageExpired(image_id) self._items.move_to_end(image_id) return stored def update_detections(self, image_id: str, detections: List[Dict[str, Any]]) -> None: with self._lock: stored = self._items.get(image_id) if stored is not None: stored.detections = list(detections) def _evict_locked(self) -> None: now = time.time() for image_id in [k for k, v in self._items.items() if now - v.created_at > self.ttl_s]: dropped = self._items.pop(image_id) self._bytes -= dropped.nbytes class RecognitionEngine: """Thread-safe facade over detector + matcher + catalog.""" def __init__( self, detector: Any, matcher: Any, inventory: Any, *, settings: Optional[EngineSettings] = None, classify_executor: Optional[ThreadPoolExecutor] = None, owns_executor: bool = True, ): self.detector = detector self.matcher = matcher self.inventory = inventory self.settings = settings or EngineSettings() workers = self.settings.classify_workers or max(2, (os.cpu_count() or 4) - 2) self.classify_workers = int(workers) self._executor = classify_executor or ThreadPoolExecutor( max_workers=self.classify_workers, thread_name_prefix="classify" ) self._owns_executor = owns_executor self._detect_lock = threading.Lock() self._inventory_lock = threading.Lock() self.image_store = ImageStore( self.settings.image_store_max, self.settings.image_store_ttl_s, self.settings.image_store_max_bytes, ) self._catalog_cache: Optional[Tuple[str, List[Dict[str, Any]]]] = None self._started_at = time.time() # ------------------------------------------------------------------ lifecycle @property def executor(self) -> ThreadPoolExecutor: return self._executor def warm_up(self) -> float: """One dummy predict so the first real frame is not the slow one.""" import numpy as np blank = np.zeros((self.settings.imgsz, self.settings.imgsz, 3), dtype="uint8") started = time.perf_counter() self.detect(blank, conf=self.settings.detect_conf) return (time.perf_counter() - started) * 1000.0 def shutdown(self) -> None: if self._owns_executor: self._executor.shutdown(wait=False, cancel_futures=True) # ------------------------------------------------------------------ images def decode_image(self, payload: bytes) -> Any: import cv2 import numpy as np if not payload: raise UndecodableImage("empty payload") array = np.frombuffer(payload, dtype=np.uint8) image = cv2.imdecode(array, cv2.IMREAD_COLOR) if image is None or image.size == 0: raise UndecodableImage("payload is not a decodable image") return image def encode_jpeg(self, image: Any, quality: int = 85) -> bytes: import cv2 ok, buffer = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality]) if not ok: return b"" return bytes(buffer) def encode_jpeg_data_url(self, image: Any, quality: int = 85) -> str: payload = self.encode_jpeg(image, quality=quality) if not payload: return "" return "data:image/jpeg;base64," + base64.b64encode(payload).decode("ascii") # ------------------------------------------------------------------ detection def detect(self, image: Any, conf: Optional[float] = None) -> List[Dict[str, Any]]: """Runs the detector under a lock; returns normalized box dicts.""" threshold = self.settings.detect_conf if conf is None else float(conf) with self._detect_lock: result = self.detector.detect( image, conf_threshold=threshold, iou_threshold=self.settings.detect_iou ) boxes = result.get("boxes") or [] normalized: List[Dict[str, Any]] = [] for index, box in enumerate(boxes): x1, y1, x2, y2 = [int(v) for v in box["box"]] if x2 <= x1 or y2 <= y1: continue normalized.append( { "id": index, "box": [x1, y1, x2, y2], "confidence": float(box.get("confidence", 0.0)), "class_id": int(box.get("class_id", 0)), "width": x2 - x1, "height": y2 - y1, "center": [(x1 + x2) // 2, (y1 + y2) // 2], } ) for index, box in enumerate(normalized): box["id"] = index return normalized def crop(self, image: Any, box: Sequence[int]) -> Any: x1, y1, x2, y2 = [int(v) for v in box] height, width = image.shape[:2] x1 = max(0, min(x1, width)) x2 = max(0, min(x2, width)) y1 = max(0, min(y1, height)) y2 = max(0, min(y2, height)) return image[y1:y2, x1:x2].copy() # ------------------------------------------------------------------ empty slots def empty_slots(self, image_shape: Sequence[int], boxes: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: """Empty-shelf gaps, computed in a 2048 px-wide reference frame so the detector's absolute pixel thresholds do not change with capture size.""" height, width = int(image_shape[0]), int(image_shape[1]) if not boxes or width <= 0 or height <= 0: return [] scale = min(MAX_UPSCALE, REFERENCE_WIDTH / float(width)) ref_boxes = [] for box in boxes: x1, y1, x2, y2 = [int(v) for v in box["box"]] ref_boxes.append( { "box": [ int(round(x1 * scale)), int(round(y1 * scale)), int(round(x2 * scale)), int(round(y2 * scale)), ] } ) ref_shape = (int(round(height * scale)), int(round(width * scale))) slots = self.detector.detect_empty_slots(ref_shape, ref_boxes) if scale == 1.0: return slots inv = 1.0 / scale for slot in slots: x1, y1, x2, y2 = slot["box"] slot["box"] = [int(round(x1 * inv)), int(round(y1 * inv)), int(round(x2 * inv)), int(round(y2 * inv))] slot["width"] = int(round(slot["width"] * inv)) slot["height"] = int(round(slot["height"] * inv)) slot["center"] = [int(round(slot["center"][0] * inv)), int(round(slot["center"][1] * inv))] return slots # ------------------------------------------------------------------ classification def score_crop(self, crop: Any, method: Optional[str] = None) -> Dict[str, float]: method = (method or self.settings.live_method).lower() if method not in CLASSIFY_METHODS: raise UnknownMethod(method) if method == "hybrid_fast": scores, _ = self.matcher.score_crop_fast(crop) else: scores, _ = self.matcher.score_crop(crop, method=method) return {str(sku): float(score) for sku, score in scores.items()} def classify_boxes( self, image: Any, boxes: Sequence[Dict[str, Any]], method: Optional[str] = None, ) -> List[Dict[str, Any]]: """Classifies every box in parallel; returns one entry per box, ordered by area.""" from identity import decide_identity method = (method or self.settings.live_method).lower() if method not in CLASSIFY_METHODS: raise UnknownMethod(method) if not boxes: return [] ordered = sorted(boxes, key=lambda b: b["width"] * b["height"], reverse=True) crops = [self.crop(image, box["box"]) for box in ordered] if method == "hybrid_fast": futures = [self._executor.submit(self.matcher.score_crop_fast, crop) for crop in crops] else: futures = [self._executor.submit(self.matcher.score_crop, crop, method) for crop in crops] results: List[Dict[str, Any]] = [] for box, future in zip(ordered, futures): try: scores, details = future.result() except Exception as exc: # a single bad crop must not fail the batch scores, details = {}, {"error": str(exc)} decision = decide_identity(scores, details=details, method=method) results.append( { "detection_id": int(box["id"]), "box": list(box["box"]), "identity": identity_payload(decision), } ) return results def inspect(self, image: Any, box: Sequence[int], method: Optional[str] = None) -> Dict[str, Any]: """Full single-crop result: decision, per-method details, crop, visualization.""" from identity import decide_identity method = (method or self.settings.live_method).lower() if method not in CLASSIFY_METHODS: raise UnknownMethod(method) crop = self.crop(image, box) if method == "hybrid_fast": scores, details = self.matcher.score_crop_fast(crop) else: scores, details = self.matcher.score_crop(crop, method=method) decision = decide_identity(scores, details=details, method=method) identity = identity_payload(decision) top_details = decision.get("details") or {} sub_details = None if isinstance(top_details, dict): raw = top_details.get("sub_details") if raw: sub_details = {k: dict(v) for k, v in raw.items()} product = None visualization = None visualized_sku = None is_candidate = False if identity["state"] == STATE_RECOGNIZED: product = self.product(identity["sku_id"]) visualized_sku = identity["sku_id"] elif identity["candidates"]: visualized_sku = identity["candidates"][0]["sku_id"] is_candidate = True if visualized_sku: try: visualization = self.matcher.generate_match_visualization( crop, visualized_sku, method=method if method != "hybrid_fast" else "hybrid" ) except Exception: visualization = None return { "box": [int(v) for v in box], "method": method, "identity": identity, "sub_details": sub_details, "product": product, "visualized_sku_id": visualized_sku, "is_candidate": is_candidate, "crop_jpeg_b64": self.encode_jpeg_data_url(crop), "match_visualization_jpeg_b64": ( self.encode_jpeg_data_url(visualization) if visualization is not None else None ), } # ------------------------------------------------------------------ catalog def catalog_version(self) -> str: digest = hashlib.sha256() 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()) except OSError: digest.update(b"missing-csv") try: 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 = [] for name, size in entries: digest.update(f"{name}:{size}".encode("utf-8")) return digest.hexdigest()[:12] def model_version(self) -> str: try: import ultralytics library = f"ultralytics-{ultralytics.__version__}" except Exception: library = "ultralytics-unknown" return f"{os.path.basename(self.settings.weights_path)}@{library}" def catalog(self) -> List[Dict[str, Any]]: version = self.catalog_version() cached = self._catalog_cache if cached and cached[0] == version: return cached[1] products: List[Dict[str, Any]] = [] for sku_id in sorted(self.inventory.list_all_skus()): product = self.inventory.get_product(sku_id) or {} products.append(self._product_payload(sku_id, product)) self._catalog_cache = (version, products) return products def product(self, sku_id: Optional[str]) -> Optional[Dict[str, Any]]: if not sku_id: return None with self._inventory_lock: product = self.inventory.get_product(sku_id) if product is None: return None payload = self._product_payload(sku_id, product) self._catalog_cache = None return payload def _product_payload(self, sku_id: str, product: Dict[str, Any]) -> Dict[str, Any]: return { "sku_id": sku_id, "product_name": product.get("product_name", ""), "price": product.get("price", 0.0), "number_in": int(product.get("number_in", 0)), "number_out": int(product.get("number_out", 0)), "number_on": int(product.get("number_on", 0)), "receipt_date": product.get("receipt_date", ""), "expiration_date": product.get("expiration_date", ""), "image_url": f"/v1/catalog/{sku_id}/image", } def catalog_image_path(self, sku_id: str) -> Optional[str]: """Resolve a template image path through the inventory map only.""" with self._inventory_lock: product = self.inventory.get_product(sku_id) if product is None: return None path = product.get("local_image_path") if not path or not os.path.exists(path): return None return path def update_stock(self, sku_id: str, delta: int) -> Optional[Dict[str, Any]]: with self._inventory_lock: ok = self.inventory.update_stock(sku_id, int(delta)) if not ok: return None product = self.inventory.get_product(sku_id) or {} payload = self._product_payload(sku_id, product) self._catalog_cache = None return payload def sku_count(self) -> int: return len(self.inventory.list_all_skus()) def uptime_s(self) -> float: return time.time() - self._started_at # ------------------------------------------------------------------ recognize def recognize( self, image: Any, *, classify: bool = True, method: Optional[str] = None, conf: Optional[float] = None, with_empty_slots: bool = True, ) -> Dict[str, Any]: """Stateless still-image recognition (photo mode and one-shot scans).""" started = time.perf_counter() boxes = self.detect(image, conf=conf) detected_at = time.perf_counter() detections: List[Dict[str, Any]] = [] identities: List[Dict[str, Any]] = [] if classify and boxes: identities = self.classify_boxes(image, boxes, method=method) identities_by_id = {row["detection_id"]: row["identity"] for row in identities} for box in boxes: detections.append( { "id": box["id"], "box": box["box"], "confidence": box["confidence"], "identity": identities_by_id.get(box["id"], identity_payload(None)), } ) slots = self.empty_slots(image.shape, boxes) if with_empty_slots else [] classified_at = time.perf_counter() summary = self.summarize(detections, slots) stored = self.image_store.put(image, detections) return { "image_id": stored.image_id, "width": stored.width, "height": stored.height, "catalog_version": self.catalog_version(), "model_version": self.model_version(), "coordinate_convention": COORDINATE_CONVENTION, "detections": detections, "empty_slots": slots, "summary": summary, "timing_ms": { "detect": round((detected_at - started) * 1000.0, 2), "classify": round((classified_at - detected_at) * 1000.0, 2), "total": round((classified_at - started) * 1000.0, 2), }, } @staticmethod def summarize( detections: Sequence[Dict[str, Any]], slots: Sequence[Dict[str, Any]], *, include_identifying: bool = False, ) -> Dict[str, Any]: counts = { STATE_RECOGNIZED: 0, STATE_NEEDS_CONFIRMATION: 0, STATE_UNRECOGNIZED: 0, STATE_DETECTED: 0, } if include_identifying: counts[STATE_IDENTIFYING] = 0 facings: Dict[str, int] = {} for row in detections: state = (row.get("identity") or {}).get("state") or STATE_DETECTED if state not in counts: state = STATE_DETECTED counts[state] += 1 sku_id = (row.get("identity") or {}).get("sku_id") if state == STATE_RECOGNIZED and sku_id: facings[sku_id] = facings.get(sku_id, 0) + 1 summary: Dict[str, Any] = { "in_view": len(detections), "recognized": counts[STATE_RECOGNIZED], "needs_confirmation": counts[STATE_NEEDS_CONFIRMATION], "unrecognized": counts[STATE_UNRECOGNIZED], "detected": counts[STATE_DETECTED], "empty_slots": len(slots), "by_sku": [ {"sku_id": sku, "facings": count} for sku, count in sorted(facings.items(), key=lambda kv: (-kv[1], kv[0])) ], } if include_identifying: summary["identifying"] = counts[STATE_IDENTIFYING] return summary def build_default_engine(settings: Optional[EngineSettings] = None) -> RecognitionEngine: """Construct the real detector/matcher/inventory. Heavy imports happen here.""" import sys settings = settings or EngineSettings() current_dir = os.path.dirname(os.path.abspath(__file__)) if current_dir not in sys.path: sys.path.insert(0, current_dir) from detector import ShelfProductDetector from feature_matcher import ProductFeatureMatcher from inventory import InventoryManager detector = ShelfProductDetector( weights_path=settings.weights_path, device=settings.device, imgsz=settings.imgsz ) matcher = ProductFeatureMatcher(stock_dir=settings.stock_dir) inventory = InventoryManager(stock_dir=settings.stock_dir, csv_filename=settings.stock_csv) return RecognitionEngine(detector, matcher, inventory, settings=settings) def settings_from_env(env: Optional[Dict[str, str]] = None) -> EngineSettings: """Build EngineSettings from DORAYAKI_* environment variables.""" env = dict(os.environ if env is None else env) base = EngineSettings() return replace( base, weights_path=env.get("DORAYAKI_WEIGHTS", base.weights_path), stock_dir=env.get("DORAYAKI_STOCK_DIR", base.stock_dir), device=env.get("DORAYAKI_DEVICE") or base.device, imgsz=int(env.get("DORAYAKI_IMGSZ", base.imgsz)), classify_workers=int(env.get("DORAYAKI_CLASSIFY_WORKERS", base.classify_workers)), live_method=env.get("DORAYAKI_LIVE_METHOD", base.live_method), live_conf=float(env.get("DORAYAKI_LIVE_CONF", base.live_conf)), live_inflight=int(env.get("DORAYAKI_LIVE_INFLIGHT", base.live_inflight)), detect_conf=float(env.get("DORAYAKI_DETECT_CONF", base.detect_conf)), ) if __name__ == "__main__": # tiny smoke check: python code/engine.py engine = build_default_engine(settings_from_env()) print("[Engine] warmup ms:", round(engine.warm_up(), 1)) print("[Engine] catalog:", engine.sku_count(), "skus, version", engine.catalog_version()) result = engine.recognize(engine.decode_image(open("image.jpg", "rb").read())) print("[Engine] detected:", len(result["detections"]), "recognized:", result["summary"]["recognized"]) print("[Engine] timing:", json.dumps(result["timing_ms"])) engine.shutdown()