From 929d77a8026cc6be912b0281fd26ffb4b9dc2c81 Mon Sep 17 00:00:00 2001 From: kiyo9w Date: Thu, 24 Sep 2026 11:52:59 +0900 Subject: [PATCH] feat: recognition engine and session tracker for multi-client serving engine.py is the process-wide, thread-safe core the server is built on: one YOLO instance behind a lock, thread-local matcher access, a classification pool so detection is never starved, the stock catalog under its own lock, empty-slot geometry normalized to a reference width (the detector's gap thresholds are absolute pixels), and a small LRU of recent still images for later classify/inspect calls. Detection output is normalized so every client sees the same box shape. tracking.py wraps ultralytics BYTETracker. The important part: ultralytics mints track ids from a process-global counter and resets it when a tracker is constructed, so with two live clients one client's ids would be re-issued while the other still held cached labels. SessionTracker gives each tracker its own counter and never touches the global one. Measured defaults (bench_engine.py, M5): detection p50 26-29 ms on CPU vs 112 ms on MPS with byte-identical boxes, so device defaults to cpu; hybrid classifies at ~51 ms/crop, 92 crops/s across 8 threads. Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recognization/code/bench_engine.py | 337 +++++++++++++++ recognization/code/engine.py | 666 +++++++++++++++++++++++++++++ recognization/code/tracking.py | 188 ++++++++ 3 files changed, 1191 insertions(+) create mode 100644 recognization/code/bench_engine.py create mode 100644 recognization/code/engine.py create mode 100644 recognization/code/tracking.py diff --git a/recognization/code/bench_engine.py b/recognization/code/bench_engine.py new file mode 100644 index 0000000..f0cf7f5 --- /dev/null +++ b/recognization/code/bench_engine.py @@ -0,0 +1,337 @@ +""" +Component benchmark for the recognition engine. + +Measures the two things that decide the live-mode defaults: how long detection +takes per frame (per device) and how long classification takes per product crop +(per method), including parallel throughput. + +Usage (from recognization/): + .venv/bin/python code/bench_engine.py + .venv/bin/python code/bench_engine.py --devices cpu mps --image image.jpg + .venv/bin/python code/bench_engine.py --compare hybrid hybrid_fast +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Optional + +current_dir = os.path.dirname(os.path.abspath(__file__)) +project_root = os.path.dirname(current_dir) +if os.path.exists(os.path.join(project_root, "stock")): + os.chdir(project_root) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) + +import cv2 # noqa: E402 +import numpy as np # noqa: E402 + + +def percentile(values: List[float], pct: float) -> float: + if not values: + return 0.0 + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int(round((pct / 100.0) * (len(ordered) - 1))))) + return ordered[index] + + +def format_table(rows: List[List[str]], headers: List[str]) -> str: + widths = [len(h) for h in headers] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(str(cell))) + lines = [" | ".join(h.ljust(widths[i]) for i, h in enumerate(headers)), + "-+-".join("-" * w for w in widths)] + for row in rows: + lines.append(" | ".join(str(cell).ljust(widths[i]) for i, cell in enumerate(row))) + return "\n".join(lines) + + +def load_image(path: str, width: Optional[int] = None) -> Any: + image = cv2.imread(path) + if image is None: + raise SystemExit(f"cannot read {path}") + if width: + scale = width / float(image.shape[1]) + image = cv2.resize(image, (int(round(image.shape[1] * scale)), int(round(image.shape[0] * scale)))) + return image + + +def bench_detection(image_path: str, devices: List[str], runs: int, imgsz: int) -> List[Dict[str, Any]]: + from detector import ShelfProductDetector + + weights = "runs/detect/train/weights/best.pt" + results: List[Dict[str, Any]] = [] + native = load_image(image_path) + frame = load_image(image_path, width=1280) + for device in devices: + try: + detector = ShelfProductDetector(weights, device=None if device == "default" else device, imgsz=imgsz) + except Exception as exc: + results.append({"device": device, "error": str(exc)}) + continue + for label, sample in (("native", native), ("1280x720-ish", frame)): + latencies: List[float] = [] + boxes = 0 + for _ in range(runs): + started = time.perf_counter() + det = detector.detect(sample) + latencies.append((time.perf_counter() - started) * 1000.0) + boxes = len(det.get("boxes", [])) + results.append( + { + "device": device, + "size": label, + "shape": f"{sample.shape[1]}x{sample.shape[0]}", + "boxes": boxes, + "p50_ms": round(statistics.median(latencies), 2), + "p95_ms": round(percentile(latencies, 95), 2), + "min_ms": round(min(latencies), 2), + } + ) + return results + + +def box_iou(a: List[int], b: List[int]) -> float: + ix1, iy1 = max(a[0], b[0]), max(a[1], b[1]) + ix2, iy2 = min(a[2], b[2]), min(a[3], b[3]) + iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1) + inter = iw * ih + if inter <= 0: + return 0.0 + area_a = (a[2] - a[0]) * (a[3] - a[1]) + area_b = (b[2] - b[0]) * (b[3] - b[1]) + return inter / max(1, area_a + area_b - inter) + + +def detection_agreement(a: List[Dict[str, Any]], b: List[Dict[str, Any]]) -> Dict[str, Any]: + if not a or not b: + return {"matched": 0, "mean_best_iou": 0.0, "count_ratio": 0.0} + ious = [] + for box_a in a: + best = max((box_iou(box_a["box"], box_b["box"]) for box_b in b), default=0.0) + ious.append(best) + return { + "matched": sum(1 for i in ious if i >= 0.5), + "mean_best_iou": round(float(np.mean(ious)), 3), + "count_ratio": round(len(a) / float(max(1, len(b))), 3), + } + + +def bench_crops(image_path: str, methods: List[str]) -> Dict[str, Any]: + from detector import ShelfProductDetector + from feature_matcher import ProductFeatureMatcher + + detector = ShelfProductDetector("runs/detect/train/weights/best.pt") + matcher = ProductFeatureMatcher(stock_dir="stock") + image = load_image(image_path) + det = detector.detect(image) + crops = [] + for box in det.get("boxes", []): + x1, y1, x2, y2 = box["box"] + crops.append(image[y1:y2, x1:x2].copy()) + report: Dict[str, Any] = {"crop_count": len(crops), "methods": {}} + for method in methods: + latencies: List[float] = [] + levels: Dict[str, int] = {} + skus: List[str] = [] + for crop in crops: + started = time.perf_counter() + result = matcher.classify_crop(crop, method=method) + latencies.append((time.perf_counter() - started) * 1000.0) + level = result.get("identity_level", "error") + levels[level] = levels.get(level, 0) + 1 + if result.get("best_sku"): + skus.append(result["best_sku"]) + report["methods"][method] = { + "p50_ms": round(statistics.median(latencies), 2) if latencies else 0.0, + "p95_ms": round(percentile(latencies, 95), 2), + "mean_ms": round(float(np.mean(latencies)), 2) if latencies else 0.0, + "levels": levels, + "accepted_skus": len(set(skus)), + } + return report + + +def bench_parallel(crops: List[Any], matcher: Any, method: str, pools: List[int], num_threads: List[int]) -> List[Dict[str, Any]]: + rows = [] + for threads in num_threads: + previous = cv2.getNumThreads() + cv2.setNumThreads(threads) + for pool in pools: + started = time.perf_counter() + with ThreadPoolExecutor(max_workers=pool) as executor: + list(executor.map(lambda c: matcher.classify_crop(c, method=method), crops)) + elapsed = time.perf_counter() - started + rows.append( + { + "cv2_threads": threads, + "pool": pool, + "crops": len(crops), + "seconds": round(elapsed, 2), + "crops_per_s": round(len(crops) / elapsed, 1), + } + ) + cv2.setNumThreads(previous) + return rows + + +def compare_methods(image_path: str, methods: List[str], extra_images: Optional[List[str]] = None) -> Dict[str, Any]: + """Agreement between two methods on the same crops (adoption gate).""" + from detector import ShelfProductDetector + from feature_matcher import ProductFeatureMatcher + + if len(methods) != 2: + raise SystemExit("--compare needs exactly two method names") + reference, candidate = methods + detector = ShelfProductDetector("runs/detect/train/weights/best.pt") + matcher = ProductFeatureMatcher(stock_dir="stock") + images = [image_path] + (extra_images or []) + accepted = 0 + agree = 0 + disagree_accepts = 0 + total = 0 + for path in images: + image = load_image(path) + for box in detector.detect(image).get("boxes", []): + x1, y1, x2, y2 = box["box"] + crop = image[y1:y2, x1:x2].copy() + if min(crop.shape[:2]) < 16: + continue + a = matcher.classify_crop(crop, method=reference) + b = matcher.classify_crop(crop, method=candidate) + total += 1 + if a.get("identity_level") == "accepted": + accepted += 1 + if a.get("best_sku") == b.get("best_sku"): + agree += 1 + elif b.get("identity_level") == "accepted": + disagree_accepts += 1 + return { + "reference": reference, + "candidate": candidate, + "crops": total, + "reference_accepted": accepted, + "agreement_on_accepted": round(agree / accepted, 4) if accepted else None, + "conflicting_accepts": disagree_accepts, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Dorayaki engine component benchmark") + parser.add_argument("--image", default="image.jpg") + parser.add_argument("--devices", nargs="*", default=["default", "cpu", "mps"]) + parser.add_argument("--runs", type=int, default=15) + parser.add_argument("--imgsz", type=int, default=640) + parser.add_argument("--methods", nargs="*", default=["surf", "orb", "brisk", "hybrid"]) + parser.add_argument("--pools", nargs="*", type=int, default=[1, 4, 8]) + parser.add_argument("--cv2-threads", nargs="*", type=int, default=[0, 1]) + parser.add_argument("--compare", nargs="*", default=None) + parser.add_argument("--compare-images", nargs="*", default=None) + parser.add_argument("--json", default=None) + args = parser.parse_args() + + print("=" * 78) + print("DORAYAKI ENGINE BENCHMARK") + print("=" * 78) + + from detector import ShelfProductDetector + + started = time.perf_counter() + detector = ShelfProductDetector("runs/detect/train/weights/best.pt") + load_ms = (time.perf_counter() - started) * 1000.0 + print(f"model load: {load_ms:.0f} ms") + + warm_image = load_image(args.image, width=1280) + started = time.perf_counter() + detector.detect(warm_image) + print(f"first predict (cold): {(time.perf_counter() - started) * 1000.0:.0f} ms") + started = time.perf_counter() + detector.detect(warm_image) + print(f"second predict (warm): {(time.perf_counter() - started) * 1000.0:.0f} ms") + + print("\n--- detection latency ---") + detection = bench_detection(args.image, args.devices, args.runs, args.imgsz) + rows = [] + for row in detection: + if "error" in row: + rows.append([row["device"], "-", "-", "-", f"ERROR: {row['error'][:40]}"]) + continue + rows.append([row["device"], row["size"], row["shape"], row["boxes"], + f"{row['p50_ms']} / {row['p95_ms']}"]) + print(format_table(rows, ["device", "size", "shape", "boxes", "p50/p95 ms"])) + + # device agreement: do mps boxes match cpu boxes? + by_device = {} + for device in args.devices: + if device in ("default", "cpu"): + continue + try: + detector_d = ShelfProductDetector("runs/detect/train/weights/best.pt", + device=None if device == "default" else device) + except Exception: + continue + by_device[device] = detector_d.detect(warm_image).get("boxes", []) + cpu_boxes = None + try: + cpu_detector = ShelfProductDetector("runs/detect/train/weights/best.pt", device="cpu") + cpu_boxes = cpu_detector.detect(warm_image).get("boxes", []) + except Exception: + pass + if cpu_boxes is not None: + for device, boxes in by_device.items(): + print(f"agreement cpu vs {device}: {json.dumps(detection_agreement(cpu_boxes, boxes))}") + + print("\n--- classification latency (real crops from", args.image, ") ---") + crop_report = bench_crops(args.image, args.methods) + print(f"crops: {crop_report['crop_count']}") + rows = [] + for method, stats in crop_report["methods"].items(): + levels = stats["levels"] + rows.append([method, f"{stats['p50_ms']} / {stats['p95_ms']}", + f"acc={levels.get('accepted', 0)}", + f"amb={levels.get('ambiguous', 0)}", + f"unk={levels.get('unknown', 0)}", + f"rej={levels.get('rejected', 0)}", + f"skus={stats['accepted_skus']}"]) + print(format_table(rows, ["method", "p50/p95 ms", "accepted", "ambiguous", "unknown", "rejected", "distinct skus"])) + + print("\n--- parallel throughput (hybrid) ---") + from feature_matcher import ProductFeatureMatcher + + matcher = ProductFeatureMatcher(stock_dir="stock") + image = load_image(args.image) + boxes = detector.detect(image).get("boxes", []) + crops = [image[b["box"][1]:b["box"][3], b["box"][0]:b["box"][2]].copy() for b in boxes] + rows = [] + for row in bench_parallel(crops, matcher, "hybrid", args.pools, args.cv2_threads): + rows.append([row["cv2_threads"], row["pool"], row["crops"], row["seconds"], row["crops_per_s"]]) + print(format_table(rows, ["cv2_threads", "pool", "crops", "seconds", "crops/s"])) + + if args.compare: + print("\n--- method agreement ---") + comparison = compare_methods(args.image, args.compare, args.compare_images) + print(json.dumps(comparison, indent=2)) + + if args.json: + with open(args.json, "w") as handle: + json.dump( + { + "model_load_ms": round(load_ms, 1), + "detection": detection, + "classification": crop_report, + }, + handle, + indent=2, + ) + print(f"\nwrote {args.json}") + + +if __name__ == "__main__": + main() diff --git a/recognization/code/engine.py b/recognization/code/engine.py new file mode 100644 index 0000000..49657b5 --- /dev/null +++ b/recognization/code/engine.py @@ -0,0 +1,666 @@ +""" +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") + + +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" + 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) + 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] + 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) + 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 + ) + 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 None, + 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)), + 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() diff --git a/recognization/code/tracking.py b/recognization/code/tracking.py new file mode 100644 index 0000000..01699ba --- /dev/null +++ b/recognization/code/tracking.py @@ -0,0 +1,188 @@ +""" +Session-scoped multi-object tracking for live recognition. + +Wraps ultralytics' BYTETracker (Kalman prediction + two-stage association) with +one critical fix: ultralytics mints track IDs from a PROCESS-GLOBAL counter +(BaseTrack._count), and BYTETracker.__init__ resets that counter. With several +live sessions that means one client's tracker creation re-issues IDs already in +use by another client, and a cached label can land on the wrong product. + +SessionTracker gives every tracker instance its own counter and never touches +the global one. The rest of the algorithm is used unmodified. + +Pinned to ultralytics 8.4.153 (see requirements.txt): this module touches +tracker internals. +""" + +from __future__ import annotations + +import itertools +import os +import sys +import threading +from typing import Any, List, Optional, Sequence + +import numpy as np + +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) + +from ultralytics.trackers.basetrack import BaseTrack, STrack +from ultralytics.trackers.byte_tracker import BYTETracker +from ultralytics.utils import IterableSimpleNamespace, yaml_load + +_BYTETRACK_YAML = os.path.join( + os.path.dirname(os.path.abspath(__import__("ultralytics").__file__)), + "cfg", "trackers", "bytetrack.yaml", +) + +DEFAULT_TRACKER_CONFIG = { + "tracker_type": "bytetrack", + "track_high_thresh": 0.25, + "track_low_thresh": 0.1, + "new_track_thresh": 0.25, + "track_buffer": 30, + "match_thresh": 0.8, + "fuse_score": True, +} + + +def _default_args() -> IterableSimpleNamespace: + try: + return IterableSimpleNamespace(**yaml_load(_BYTETRACK_YAML)) + except Exception: + return IterableSimpleNamespace(**DEFAULT_TRACKER_CONFIG) + + +class _Detections: + """Duck-typed stand-in for ultralytics Results: only what BYTETracker reads. + + Must NOT expose `xywhr`; the tracker uses its presence to pick rotated-box + handling. + """ + + __slots__ = ("xywh", "conf", "cls") + + def __init__(self, boxes_xyxy: np.ndarray, scores: np.ndarray): + boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4) + scores = np.asarray(scores, dtype=np.float32).reshape(-1) + widths = boxes[:, 2] - boxes[:, 0] + heights = boxes[:, 3] - boxes[:, 1] + self.xywh = np.stack( + [boxes[:, 0] + widths / 2.0, boxes[:, 1] + heights / 2.0, widths, heights], axis=1 + ).astype(np.float32) + self.conf = scores + self.cls = np.zeros(len(scores), dtype=np.float32) + + def __len__(self) -> int: + return len(self.conf) + + def __getitem__(self, mask: Any) -> "_Detections": + mask = np.asarray(mask) + subset = _Detections.__new__(_Detections) + subset.xywh = self.xywh[mask] + subset.conf = self.conf[mask] + subset.cls = self.cls[mask] + return subset + + +class _SessionByteTrack(BYTETracker): + """BYTETracker with a per-instance ID counter.""" + + def __init__(self, args: IterableSimpleNamespace): + self._id_counter = itertools.count(1) + self._id_lock = threading.Lock() + # super().__init__ calls self.reset_id(); the override below only touches + # this instance's counter, so the global one is never disturbed. + super().__init__(args) + self.track_class = self._make_track_class() + + def _next_id(self) -> int: + with self._id_lock: + return next(self._id_counter) + + def _make_track_class(self): + tracker = self + + class SessionSTrack(STrack): + @staticmethod + def next_id() -> int: + return tracker._next_id() + + SessionSTrack.__name__ = "SessionSTrack" + return SessionSTrack + + def reset_id(self) -> None: # type: ignore[override] + """Reset only this tracker's counter; never BaseTrack._count.""" + with self._id_lock: + self._id_counter = itertools.count(1) + self.track_class = self._make_track_class() + + +class TrackOut: + """One tracked object for a frame.""" + + __slots__ = ("track_id", "box", "score", "det_index") + + def __init__(self, track_id: int, box: List[int], score: float, det_index: Optional[int]): + self.track_id = int(track_id) + self.box = [int(v) for v in box] + self.score = float(score) + self.det_index = None if det_index is None or det_index < 0 else int(det_index) + + def __repr__(self) -> str: # pragma: no cover - debug helper + return f"TrackOut(id={self.track_id}, box={self.box}, score={self.score:.3f}, det={self.det_index})" + + +class SessionTracker: + """Per-session tracker. Used by one session at a time, under its lock.""" + + def __init__( + self, + *, + high: float = 0.25, + low: float = 0.1, + new: float = 0.25, + buffer: int = 30, + match: float = 0.8, + fuse: bool = True, + ): + args = _default_args() + args.track_high_thresh = float(high) + args.track_low_thresh = float(low) + args.new_track_thresh = float(new) + args.track_buffer = int(buffer) + args.match_thresh = float(match) + args.fuse_score = bool(fuse) + self.args = args + self._tracker = _SessionByteTrack(args) + + def update( + self, + boxes_xyxy: Sequence[Sequence[int]], + scores: Sequence[float], + ) -> List[TrackOut]: + boxes = np.asarray(boxes_xyxy, dtype=np.float32).reshape(-1, 4) + scores_array = np.asarray(scores, dtype=np.float32).reshape(-1) + if len(boxes) != len(scores_array): + raise ValueError("boxes and scores must have the same length") + detections = _Detections(boxes, scores_array) + rows = self._tracker.update(detections) + outputs: List[TrackOut] = [] + if rows is None or len(rows) == 0: + return outputs + for row in np.asarray(rows): + # [x1, y1, x2, y2, track_id, score, cls, idx] + outputs.append( + TrackOut( + track_id=int(row[4]), + box=[int(round(row[0])), int(round(row[1])), int(round(row[2])), int(round(row[3]))], + score=float(row[5]), + det_index=int(row[7]) if len(row) > 7 else None, + ) + ) + return outputs + + def reset(self) -> None: + self._tracker.reset()