""" Dorayaki recognition server (どら焼き). The computer-vision model runs here, as a service. Capture clients talk to it over a versioned REST API (stills, catalog, inspect) and a WebSocket stream (live frames in, tracked products and labels out). Any number of clients can connect at once; no client shares state with another. Run: python code/app.py --port 8080 python code/app.py --host 0.0.0.0 --port 8080 # LAN, set DORAYAKI_API_TOKEN python code/app.py --ssl-certfile cert.pem --ssl-keyfile key.pem The web client is served from / and is a pure client of the API; it can also be pointed at another server with ?server=https://host:port. """ from __future__ import annotations import argparse import asyncio import base64 import hmac import json import logging import os import re import sys import threading import time import uuid 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") os.environ.setdefault("YOLO_AUTOINSTALL", "false") 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) from fastapi import ( # noqa: E402 Depends, FastAPI, Header, HTTPException, Request, WebSocket, WebSocketDisconnect, ) from fastapi.exceptions import RequestValidationError # noqa: E402 from fastapi.middleware.cors import CORSMiddleware # noqa: E402 from fastapi.responses import FileResponse, JSONResponse, Response # noqa: E402 from fastapi.staticfiles import StaticFiles # noqa: E402 from pydantic import BaseModel, Field # noqa: E402 from engine import ( # noqa: E402 CLASSIFY_METHODS, COORDINATE_CONVENTION, EngineSettings, ImageExpired, RecognitionEngine, STATE_DETECTED, UndecodableImage, UnknownMethod, build_default_engine, identity_payload, settings_from_env, ) from demo import DemoCamera, DemoSourceError # noqa: E402 from record import FrameRecorder # noqa: E402 from live import LiveSession, LiveSettings # noqa: E402 LOGGER = logging.getLogger("dorayaki") MAX_UPLOAD_BYTES = 20 * 1024 * 1024 MAX_FRAME_BYTES = 4 * 1024 * 1024 DEFAULT_MAX_SESSIONS = 4 WS_IDLE_TIMEOUT_S = 120.0 CATALOG_VERSION_FALLBACK = "unknown" # --------------------------------------------------------------------------- # settings # --------------------------------------------------------------------------- @dataclass class ServerSettings: host: str = "127.0.0.1" port: int = 8080 token: Optional[str] = None cors_origins: str = "*" max_sessions: int = DEFAULT_MAX_SESSIONS live_method: str = "hybrid" live_conf: float = 0.1 live_tags: bool = True demo_walk: bool = False # A live session can be recorded so it becomes data to label (record.py). record_frames: Optional[str] = None record_every: int = 3 log_level: str = "info" ssl_certfile: Optional[str] = None ssl_keyfile: Optional[str] = None def server_settings_from_env() -> ServerSettings: base = ServerSettings() return ServerSettings( host=os.environ.get("DORAYAKI_HOST", base.host), port=int(os.environ.get("DORAYAKI_PORT", base.port)), token=os.environ.get("DORAYAKI_API_TOKEN") or None, cors_origins=os.environ.get("DORAYAKI_CORS_ORIGINS", base.cors_origins), max_sessions=int(os.environ.get("DORAYAKI_MAX_SESSIONS", base.max_sessions)), live_method=os.environ.get("DORAYAKI_LIVE_METHOD", base.live_method), live_conf=float(os.environ.get("DORAYAKI_LIVE_CONF", base.live_conf)), live_tags=os.environ.get("DORAYAKI_LIVE_TAGS", "1").lower() not in ("0", "false", "no"), log_level=os.environ.get("DORAYAKI_LOG_LEVEL", base.log_level), ssl_certfile=os.environ.get("DORAYAKI_SSL_CERTFILE") or None, ssl_keyfile=os.environ.get("DORAYAKI_SSL_KEYFILE") or None, ) # --------------------------------------------------------------------------- # request models # --------------------------------------------------------------------------- class StockUpdateBody(BaseModel): delta: int = Field(..., ge=-100, le=100, description="units added (+) or dispensed (-)") class ClassifyBody(BaseModel): method: Optional[str] = None detection_ids: Optional[List[int]] = None class InspectBody(BaseModel): box: List[int] = Field(..., min_length=4, max_length=4) method: Optional[str] = None # --------------------------------------------------------------------------- # helpers # --------------------------------------------------------------------------- def _error(code: str, message: str, status_code: int) -> HTTPException: return HTTPException(status_code=status_code, detail={"error": {"code": code, "message": message}}) def _require_method(method: Optional[str]) -> str: chosen = (method or "hybrid").lower() if chosen not in CLASSIFY_METHODS: raise _error("BAD_REQUEST", f"unknown method '{method}'", 400) return chosen def get_available_shelf_images(add_dir: str = "add") -> List[Dict[str, str]]: """Same demo set the old /api/shelf_samples returned, now id-keyed.""" items: List[Dict[str, str]] = [] if os.path.exists("image.jpg"): items.append({"id": "main_shelf", "path": "image.jpg", "label": "image.jpg (Main Shelf)"}) if os.path.exists(add_dir): valid_exts = (".webp", ".jpg", ".jpeg", ".png") files = sorted( [f for f in os.listdir(add_dir) if f.lower().endswith(valid_exts)], key=lambda x: (len(x), x), ) for name in files: items.append( { "id": name.rsplit(".", 1)[0], "path": os.path.join(add_dir, name).replace("\\", "/"), "label": f"{add_dir}/{name}", } ) mosaic = "runs/detect/train/train_batch0.jpg" if os.path.exists(mosaic): items.append({"id": "train_mosaic", "path": mosaic, "label": "train_batch0.jpg (Mosaic)"}) return items def build_legacy_scans(engine: RecognitionEngine): """The /api/sessions* routes kept from the earlier session model, now backed by the real engine instead of the fake adapter.""" from observe import PerceptionService, ScanService from yolo_adapter import YoloMatcherAdapter class _DetectorFacade: def __init__(self, engine: RecognitionEngine): self.engine = engine self.weights_path = engine.settings.weights_path def detect(self, image, conf_threshold: float = 0.25, iou_threshold: float = 0.45): boxes = self.engine.detect(image, conf=conf_threshold) return {"status": "success", "boxes": boxes, "total_detected": len(boxes)} def detect_empty_slots(self, image_shape, boxes, min_gap_ratio: float = 0.75): return self.engine.empty_slots(image_shape, boxes) def crop_roi(self, image, box): return self.engine.crop(image, box) class _MatcherFacade: def __init__(self, engine: RecognitionEngine): self.engine = engine def classify_crop(self, crop, method: str = "hybrid"): scores = self.engine.score_crop(crop, method=method) from identity import decide_identity return decide_identity(scores, method=method) adapter = YoloMatcherAdapter(_DetectorFacade(engine), _MatcherFacade(engine), method="hybrid") return ScanService(PerceptionService(adapter)) # --------------------------------------------------------------------------- # app factory # --------------------------------------------------------------------------- def create_app( engine: RecognitionEngine, settings: Optional[ServerSettings] = None, *, web_dir: Optional[str] = None, scans: Any = None, ) -> FastAPI: settings = settings or server_settings_from_env() app = FastAPI( title="Dorayaki Recognition Server", version="1.0", description="Shelf product detection, tracking and classification for capture clients.", ) app.add_middleware( CORSMiddleware, allow_origins=[o.strip() for o in settings.cors_origins.split(",") if o.strip()] or ["*"], allow_credentials=False, allow_methods=["*"], allow_headers=["*"], ) scans = scans if scans is not None else build_legacy_scans(engine) samples = {item["id"]: item for item in get_available_shelf_images("add")} demo_camera = DemoCamera(samples) demo_cameras: Dict[float, DemoCamera] = {} sessions: Dict[str, LiveSession] = {} sessions_lock = asyncio.Lock() app.state.engine = engine app.state.sessions = sessions async def authorize(authorization: Optional[str] = Header(default=None)) -> None: if not settings.token: return expected = f"Bearer {settings.token}" if not authorization or not hmac.compare_digest(authorization, expected): raise _error("UNAUTHORIZED", "missing or invalid bearer token", 401) # ---------------------------------------------------------------- health @app.get("/v1/health") async def health() -> Dict[str, Any]: return { "status": "ok", "model": { "version": engine.model_version(), "device": getattr(engine.settings, "device", None) or "auto", "imgsz": getattr(engine.settings, "imgsz", 640), "warm": True, }, "catalog": {"version": engine.catalog_version(), "sku_count": engine.sku_count()}, "identity": { "backend": getattr(engine.settings, "identity_backend", "matcher"), "gallery_views": ( engine.identity_gallery().view_count() if getattr(engine.settings, "identity_backend", "") == "gallery" and engine.identity_gallery() else None ), }, "ocr": { # Without this an operator could not tell a broken OCR from a shelf # with no readable tags. "engine": getattr(engine.ocr_engine, "name", "none"), "available": bool(getattr(engine.ocr_engine, "available", lambda: False)()), }, "recording": recorder.stats() if recorder is not None else None, "sessions_active": len(sessions), "uptime_s": round(engine.uptime_s(), 1), } # ---------------------------------------------------------------- catalog @app.get("/v1/catalog", dependencies=[Depends(authorize)]) async def catalog() -> Dict[str, Any]: return {"catalog_version": engine.catalog_version(), "products": engine.catalog()} @app.get("/v1/catalog/{sku_id}/image", dependencies=[Depends(authorize)]) async def catalog_image(sku_id: str) -> Response: path = engine.catalog_image_path(sku_id) if not path: raise _error("NOT_FOUND", f"no catalog image for '{sku_id}'", 404) return FileResponse(path) @app.post("/v1/catalog/{sku_id}/stock", dependencies=[Depends(authorize)]) async def update_stock(sku_id: str, body: StockUpdateBody) -> Dict[str, Any]: if body.delta == 0: raise _error("VALIDATION_ERROR", "delta must not be zero", 422) product = engine.update_stock(sku_id, body.delta) if product is None: raise _error("NOT_FOUND", f"unknown sku '{sku_id}'", 404) return {"status": "success", "product": product} # ---------------------------------------------------------------- samples @app.get("/v1/samples", dependencies=[Depends(authorize)]) async def list_samples() -> Dict[str, Any]: listing = [] for item in samples.values(): entry = {"id": item["id"], "label": item["label"], "url": f"/v1/samples/{item['id']}"} try: # size is a convenience for clients that plan a demo walk width, height = demo_camera._size(item["id"]) entry["width"], entry["height"] = width, height except DemoSourceError: pass listing.append(entry) return {"samples": listing} # ---------------------------------------------------------------- demo camera # The demo replay lives in its own module: it replaces the live tracks with a # still recognition mapped into each walk frame, which is a film-making device, # not a perception result. Every payload it produces is marked synthetic. demo_recognition = None recorder = FrameRecorder(settings.record_frames, every=settings.record_every) \ if settings.record_frames else None def _demo_samples(raw: Optional[str]) -> List[str]: if raw: wanted = [part.strip() for part in raw.split(",") if part.strip()] unknown = [part for part in wanted if part not in samples] if unknown: raise _error("NOT_FOUND", f"unknown sample(s): {', '.join(unknown)}", 404) return wanted # main_shelf first: its products are the ones the catalog covers, so a # demo walk recognises real items instead of unknown shelf neighbours. for preferred in ("main_shelf", "image6"): if preferred in samples: return [preferred] return [next(iter(samples))] if samples else [] # Which sample the walk uses, resolved once. The replay object itself is built # later, because it needs the still tag pass. demo_sample_ids: List[str] = [] if settings.demo_walk: try: demo_sample_ids = _demo_samples(None) demo_camera.warm(demo_sample_ids) except Exception as exc: LOGGER.warning("demo walk unavailable: %s", exc) def _demo_window(width: Optional[int], height: Optional[int]) -> tuple: win_w = int(width or 960) win_h = int(height or 540) if not (160 <= win_w <= 2048 and 120 <= win_h <= 2048): raise _error("BAD_REQUEST", "window must be between 160x120 and 2048x2048", 400) return win_w, win_h def _demo_camera(pan: Optional[float]) -> DemoCamera: """One camera per pacing value (the walk geometry depends on it).""" if pan is None: return demo_camera key = round(float(pan), 4) camera = demo_cameras.get(key) if camera is None: camera = DemoCamera(samples, pan_overlap=key) demo_cameras[key] = camera return camera @app.get("/v1/demo/plan", dependencies=[Depends(authorize)]) async def demo_plan(samples_param: Optional[str] = None, width: Optional[int] = None, height: Optional[int] = None, pan: Optional[float] = None) -> Dict[str, Any]: """Geometry of the virtual walk (demo-only capture source).""" try: ids = _demo_samples(samples_param) if not ids: raise _error("NOT_FOUND", "no samples available", 404) camera = _demo_camera(pan) return {"status": "success", **camera.plan(ids, _demo_window(width, height))} except DemoSourceError as exc: raise _error("NOT_FOUND", str(exc), 404) @app.get("/v1/demo/frame", dependencies=[Depends(authorize)]) async def demo_frame(step: int = 0, samples_param: Optional[str] = None, width: Optional[int] = None, height: Optional[int] = None, quality: int = 65, pan: Optional[float] = None) -> Response: """One JPEG frame of the virtual walk: a store employee sweeping the shelf. Demo-only. The client advances `step` (it wraps), so the walk's speed follows whatever frame rate the client can sustain. """ try: ids = _demo_samples(samples_param) if not ids: raise _error("NOT_FOUND", "no samples available", 404) payload, position = _demo_camera(pan).frame( ids, step, _demo_window(width, height), max(30, min(95, int(quality))) ) except DemoSourceError as exc: raise _error("NOT_FOUND", str(exc), 404) return Response( content=payload, media_type="image/jpeg", headers={ "X-Demo-Sample": str(position["sample"]), "X-Demo-Bay": str(position["bay_index"] + 1), "X-Demo-Step": str(position["step"]), "X-Demo-Steps": str(position["steps"]), "X-Demo-X": str(position["x"]), "X-Demo-Y": str(position["y"]), }, ) @app.get("/v1/samples/{sample_id}", dependencies=[Depends(authorize)]) async def get_sample(sample_id: str) -> Response: item = samples.get(sample_id) if not item or not os.path.exists(item["path"]): raise _error("NOT_FOUND", f"no sample '{sample_id}'", 404) return FileResponse(item["path"]) # ---------------------------------------------------------------- recognize @app.post("/v1/recognize", dependencies=[Depends(authorize)]) async def recognize( request: Request, classify: bool = True, method: Optional[str] = None, conf: Optional[float] = None, empty_slots: bool = True, tags: bool = False, ) -> Dict[str, Any]: if method is not None: _require_method(method) content_type = (request.headers.get("content-type") or "").split(";")[0].strip().lower() if content_type == "multipart/form-data": form = await request.form() upload = form.get("file") if upload is None or not hasattr(upload, "read"): raise _error("BAD_REQUEST", "multipart body needs a 'file' field", 400) payload = await upload.read() else: if content_type and not content_type.startswith("image/") and content_type != "application/octet-stream": raise _error("UNSUPPORTED_MEDIA_TYPE", f"unsupported content type '{content_type}'", 415) payload = await request.body() if not payload: raise _error("BAD_REQUEST", "empty request body", 400) if len(payload) > MAX_UPLOAD_BYTES: raise _error("PAYLOAD_TOO_LARGE", f"image larger than {MAX_UPLOAD_BYTES} bytes", 413) try: image = engine.decode_image(payload) except UndecodableImage as exc: raise _error("UNDECODABLE_IMAGE", str(exc), 422) try: result = engine.recognize( image, classify=classify, method=method, conf=conf, with_empty_slots=empty_slots, ) except UnknownMethod as exc: raise _error("BAD_REQUEST", str(exc), 400) if tags: result = _attach_still_tags(result, image) return result def _is_price_card(crop: Any) -> bool: """Is this detection a shelf price card rather than a product? Measured on a supermarket photo whose shelf edge carries big price cards: the detector boxes them as products, because a card is roughly the shape of a pack. Their text gives them away - a price card's characters are 30-44 % digits and it prints the tax-excluded price (本体), while a product's are 3-19 % digits. Reporting a price card as a product would put a box labelled with a product name on a price tag, which is worse than not detecting it. """ ocr_engine = getattr(engine, "ocr_engine", None) if ocr_engine is None or not getattr(ocr_engine, "available", lambda: False)(): return False try: import cv2 big = cv2.resize(crop, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC) lines = ocr_engine.read(big).lines except Exception: return False text = " ".join(line.text for line in lines) compact = re.sub(r"\s", "", text) if not compact: return False digits = sum(1 for char in compact if char.isdigit() or char in "¥￥./") / len(compact) return ("本体" in text and digits >= 0.20) or digits >= 0.30 def _drop_price_cards(result: Dict[str, Any], image: Any) -> int: """Remove detections that are price cards, and say how many were removed.""" kept, dropped = [], 0 for detection in result["detections"]: crop = engine.crop(image, detection["box"]) if crop is not None and crop.size and _is_price_card(crop): dropped += 1 continue kept.append(detection) if dropped: result["detections"] = kept result["summary"] = engine.summarize(kept, result.get("empty_slots") or []) result["price_cards_dropped"] = dropped return dropped def _attach_still_tags(result: Dict[str, Any], image: Any) -> Dict[str, Any]: """Run one synchronous tag pass over a still and merge its evidence.""" ocr_engine = getattr(engine, "ocr_engine", None) if ocr_engine is None or not getattr(ocr_engine, "available", lambda: False)(): result["price_tags"] = {"readable": 0, "unreadable": 0, "absent": 0} result["compliance"] = {"violations": 0, "by_rule": {}, "rows": [], "hints": []} result["ocr_engine"] = "none" return result from reconcile import attach_tags_to_boxes from tags import TagService, TagSettings # Before anything else: a price card is not a product. _drop_price_cards(result, image) boxes = [ {"id": detection["id"], "box": detection["box"]} for detection in result["detections"] ] identities = { int(detection["id"]): detection["identity"] for detection in result["detections"] } catalog_by_sku = {product["sku_id"]: product for product in engine.catalog()} service = TagService(ocr_engine, settings=TagSettings()) started = time.perf_counter() bundle = attach_tags_to_boxes(service, image, boxes, catalog_by_sku, identities) elapsed_ms = (time.perf_counter() - started) * 1000.0 # Tag-first identity: the store's own label, read from the tag printed # under the facings it names. It fills identities the packaging pass could # not resolve and never overwrites one it did. tag_identities = bundle.get("tag_identities") or {} for detection in result["detections"]: derived = tag_identities.get(int(detection["id"])) if not derived: continue # The tag's claim is always reported, so a client can compare the two # identity channels; it only *fills* the identity when the packaging # pass could not name the product. detection["tag_identity"] = dict(derived) if (detection.get("identity") or {}).get("state") == "recognized" \ and (detection["identity"].get("sku_id") == derived["sku_id"] or detection["identity"].get("source") == "tag_name"): continue product = catalog_by_sku.get(derived["sku_id"]) or {} detection["identity"] = { "state": "recognized", "sku_id": derived["sku_id"], "identity_level": "tag_derived", "confidence": derived["score"], "score": derived["score"], "source": "tag_name", "candidates": [], "attempts": 1, "method": "tag_name", } detection["product"] = { "sku_id": product.get("sku_id"), "product_name": product.get("product_name"), "price": product.get("price"), } for detection in result["detections"]: box_id = int(detection["id"]) tag = bundle["by_box"].get(box_id) if tag is not None: detection["price_tag"] = tag reconciliation = bundle["reconciliations"].get(box_id) if reconciliation is not None: detection["reconciliation"] = reconciliation result["price_tags"] = bundle["counts"] result["compliance"] = bundle["compliance"] result["ocr_engine"] = ocr_engine.name result["tag_timing_ms"] = round(elapsed_ms, 1) engine.image_store.set_tag_bundle(result["image_id"], bundle) return result if settings.demo_walk and demo_sample_ids: # Built here because the replay needs the still tag pass, and it recognises # the whole shelf once up front so the walk's first frame is already # labelled. Every payload it produces is marked synthetic. try: from demo_replay import DemoRecognition demo_recognition = DemoRecognition(engine, demo_camera, samples, _attach_still_tags) demo_recognition.warm(demo_sample_ids[0]) LOGGER.info("demo walk: '%s' recognised up front", demo_sample_ids[0]) except Exception as exc: LOGGER.warning("demo walk warm-up failed: %s", exc) @app.post("/v1/images/{image_id}/classify", dependencies=[Depends(authorize)]) async def classify_image(image_id: str, body: ClassifyBody) -> Dict[str, Any]: method = _require_method(body.method) try: stored = engine.image_store.get(image_id) except ImageExpired: raise _error("IMAGE_EXPIRED", f"image '{image_id}' is no longer available", 404) detections = stored.detections if body.detection_ids is not None: wanted = set(int(i) for i in body.detection_ids) detections = [d for d in detections if int(d["id"]) in wanted] started = time.perf_counter() rows = engine.classify_boxes(stored.image, detections, method=method) rows_by_id = {row["detection_id"]: row for row in rows} for detection in stored.detections: row = rows_by_id.get(int(detection["id"])) if row is not None: detection["identity"] = row["identity"] engine.image_store.update_detections(image_id, stored.detections) # Tags were read before the identities existed; reconcile them now. compliance = None if stored.tag_bundle: from reconcile import reconcile_bundle catalog_by_sku = {product["sku_id"]: product for product in engine.catalog()} refreshed = reconcile_bundle(stored.tag_bundle, stored.detections, catalog_by_sku) for detection in stored.detections: decision = refreshed["reconciliations"].get(int(detection["id"])) if decision is not None: detection["reconciliation"] = decision compliance = refreshed["compliance"] identities = [] for detection in detections: row = { "detection_id": int(detection["id"]), "box": list(detection["box"]), "identity": detection["identity"], } # Re-reconciled after classification; without it a client would keep # the `tag_only` verdict it got when identities did not exist yet. if detection.get("reconciliation") is not None: row["reconciliation"] = detection["reconciliation"] if detection.get("price_tag") is not None: row["price_tag"] = detection["price_tag"] identities.append(row) summary = engine.summarize(stored.detections, []) payload = { "image_id": image_id, "method": method, "identities": identities, "summary": summary, "timing_ms": {"classify": round((time.perf_counter() - started) * 1000.0, 2)}, } if compliance is not None: payload["compliance"] = compliance return payload @app.post("/v1/images/{image_id}/inspect", dependencies=[Depends(authorize)]) async def inspect_image(image_id: str, body: InspectBody) -> Dict[str, Any]: method = _require_method(body.method) try: stored = engine.image_store.get(image_id) except ImageExpired: raise _error("IMAGE_EXPIRED", f"image '{image_id}' is no longer available", 404) return engine.inspect(stored.image, body.box, method=method) # ---------------------------------------------------------------- live stream @app.websocket("/v1/stream") async def stream(websocket: WebSocket) -> None: if settings.token: supplied = websocket.query_params.get("token") or "" header = websocket.headers.get("authorization") or "" expected_header = f"Bearer {settings.token}" if not ( hmac.compare_digest(supplied, settings.token) or hmac.compare_digest(header, expected_header) ): await websocket.close(code=1008) return async with sessions_lock: if len(sessions) >= settings.max_sessions: await websocket.close(code=1013) return session_id = str(uuid.uuid4()) session = LiveSession( session_id, engine, settings=LiveSettings( method=settings.live_method, conf=settings.live_conf, max_inflight_per_session=engine.settings.live_inflight, max_attempts=engine.settings.live_attempts, tags=settings.live_tags, ), ) session.refresh_catalog() if settings.live_tags: session.enable_tags() sessions[session_id] = session await websocket.accept() send_lock = asyncio.Lock() loop = asyncio.get_running_loop() demo_steps = [0] async def send(payload: Dict[str, Any]) -> None: async with send_lock: await websocket.send_text(json.dumps(payload)) await send( { "type": "hello", "protocol": 1, "session_id": session_id, "catalog_version": engine.catalog_version(), "model": { "version": engine.model_version(), "device": getattr(engine.settings, "device", None) or "auto", "imgsz": getattr(engine.settings, "imgsz", 640), }, "config": session.config(), "limits": {"max_frame_bytes": MAX_FRAME_BYTES}, } ) async def processing_loop() -> None: while not session.closed: pending = await loop.run_in_executor(None, session.take_pending, 0.25) if pending is None: continue frame_id, payload = pending try: result = await loop.run_in_executor(None, session.process_frame, frame_id, payload) except UndecodableImage: await send({"type": "error", "code": "UNDECODABLE_IMAGE", "frame_id": frame_id, "message": "frame is not a readable image"}) continue except Exception as exc: # keep the stream alive on one bad frame LOGGER.exception("live frame failed: %s", exc) await send({"type": "error", "code": "INTERNAL", "frame_id": frame_id, "message": "frame processing failed"}) continue if recorder is not None: # Record what the client sent and what we saw in it, so the walk # can be labelled later instead of only measured on photographs. recorder.write(session_id, frame_id, payload, result) if demo_recognition is not None and demo_sample_ids: # The walk's frames are served from the shelf's full recognition, # so every product is already labelled with name and price. sample_id = demo_sample_ids[0] step = demo_steps[0] try: result = await loop.run_in_executor( None, demo_recognition.apply, result, sample_id, step ) demo_steps[0] = (step + 1) % max(1, demo_recognition.steps(sample_id)) except Exception as exc: # a demo failure must not kill the stream LOGGER.exception("demo walk frame failed: %s", exc) await send(result) async def handle_inspect(message: Dict[str, Any]) -> None: request_id = message.get("request_id") try: track_id = int(message.get("track_id")) except (TypeError, ValueError): await send({"type": "error", "code": "BAD_MESSAGE", "request_id": request_id, "message": "track_id required"}) return method = message.get("method") if method is not None: try: method = _require_method(method) except HTTPException as exc: await send({"type": "error", "code": "BAD_MESSAGE", "request_id": request_id, "message": str(exc.detail)}) return try: result = await loop.run_in_executor(None, session.inspect, track_id, method) except KeyError: await send({"type": "error", "code": "TRACK_NOT_IN_VIEW", "request_id": request_id, "message": "track is not in the latest frame"}) return except Exception as exc: LOGGER.exception("inspect failed: %s", exc) await send({"type": "error", "code": "INTERNAL", "request_id": request_id, "message": "inspect failed"}) return payload = {"type": "inspect_result", "request_id": request_id, "track_id": track_id} payload.update(result) await send(payload) processing = asyncio.create_task(processing_loop()) async def handle_control(message: Dict[str, Any]) -> None: kind = message.get("type") if kind == "config": try: config = session.configure( method=message.get("method"), conf=message.get("conf"), empty_slots=message.get("empty_slots"), tags=message.get("tags"), ) except (TypeError, ValueError): await send({"type": "error", "code": "BAD_MESSAGE", "message": "bad config values"}) return await send({"type": "config_ack", "config": config}) elif kind == "reset": session.reset() await send({"type": "reset_ack"}) elif kind == "ping": await send({"type": "pong", "t": message.get("t")}) elif kind == "inspect": await handle_inspect(message) else: await send({"type": "error", "code": "BAD_MESSAGE", "message": f"unknown type '{kind}'"}) try: while True: try: message = await asyncio.wait_for(websocket.receive(), timeout=WS_IDLE_TIMEOUT_S) except asyncio.TimeoutError: break if message.get("type") == "websocket.disconnect": break text = message.get("text") data = message.get("bytes") if data is not None: if len(data) < 5: await send({"type": "error", "code": "BAD_FRAME", "message": "frame shorter than 5 bytes"}) continue if len(data) > MAX_FRAME_BYTES + 4: await send({"type": "error", "code": "PAYLOAD_TOO_LARGE", "message": "frame too large"}) continue frame_id = int.from_bytes(data[:4], "big") session.submit_frame(frame_id, data[4:]) elif text is not None: try: parsed = json.loads(text) if not isinstance(parsed, dict): raise ValueError("not an object") except (json.JSONDecodeError, ValueError): await send({"type": "error", "code": "BAD_MESSAGE", "message": "invalid JSON"}) continue await handle_control(parsed) except WebSocketDisconnect: pass finally: session.close() processing.cancel() try: await processing except (asyncio.CancelledError, Exception): pass async with sessions_lock: sessions.pop(session_id, None) # ---------------------------------------------------------------- legacy session routes @app.post("/api/sessions") async def legacy_start_session(payload: Dict[str, Any] = None) -> Dict[str, Any]: payload = payload or {} session = scans.start_session( zone_id=str(payload.get("zone_id") or "zone-1"), catalog_version=str(payload.get("catalog_version") or "stock-v1"), capture_kind=str(payload.get("capture_kind") or "sampled"), ) return {"status": "success", "session": session.snapshot()} @app.post("/api/sessions/{session_id}/frames") async def legacy_observe_frame(session_id: str, payload: Dict[str, Any] = None) -> Dict[str, Any]: from observe import FrameRequest, ObserveError payload = payload or {} try: session = scans.get_session(session_id) raw = payload.get("frame_bytes_b64") or payload.get("image_base64") or "" if "," in raw: raw = raw.split(",", 1)[1] request = FrameRequest( frame_id=str(payload.get("frame_id") or ""), frame_bytes=base64.b64decode(raw) if raw else b"", sequence=int(payload.get("sequence", 0)), captured_at=float(payload.get("captured_at") or 0), session_id=session.session_id, zone_id=session.zone_id, catalog_version=session.catalog_version, capture_kind=str(payload.get("capture_kind") or session.capture_kind), width=payload.get("width"), height=payload.get("height"), ) if not request.frame_id: raise _error("BAD_REQUEST", "frame_id required", 400) observation = scans.observe_frame(request) return {"status": "success", "observation": observation.as_dict(), "session": session.snapshot()} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) @app.get("/api/sessions/{session_id}") async def legacy_get_session(session_id: str) -> Dict[str, Any]: from observe import ObserveError try: return {"status": "success", "session": scans.get_session(session_id).snapshot()} except ObserveError as exc: raise _error("NOT_FOUND", str(exc), 404) @app.post("/api/sessions/{session_id}/review") async def legacy_review(session_id: str, payload: Dict[str, Any] = None) -> Dict[str, Any]: from observe import ObserveError payload = payload or {} try: snapshot = scans.review( session_id, actor=str(payload.get("actor") or "operator"), reason=str(payload.get("reason") or "manual-review"), note=payload.get("note"), ) return {"status": "success", "session": snapshot} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) @app.post("/api/sessions/{session_id}/finalize") async def legacy_finalize(session_id: str) -> Dict[str, Any]: from observe import ObserveError try: return {"status": "success", "session": scans.finalize(session_id)} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) @app.post("/api/sessions/{session_id}/cancel") async def legacy_cancel(session_id: str) -> Dict[str, Any]: from observe import ObserveError try: return {"status": "success", "session": scans.cancel(session_id)} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) @app.post("/api/sessions/{session_id}/rescan") async def legacy_rescan(session_id: str) -> Dict[str, Any]: from observe import ObserveError try: return {"status": "success", "session": scans.rescan(session_id).snapshot()} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) @app.post("/api/compare_live_vs_still") async def legacy_compare(payload: Dict[str, Any] = None) -> Dict[str, Any]: from observe import ObserveError payload = payload or {} try: comparison = scans.compare_live_vs_still( str(payload.get("sampled_session_id")), str(payload.get("still_session_id")) ) return {"status": "success", "comparison": comparison} except ObserveError as exc: raise _error("CONFLICT", str(exc), 409) # ---------------------------------------------------------------- errors @app.exception_handler(RequestValidationError) async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse: first = (exc.errors() or [{}])[0] location = ".".join(str(part) for part in first.get("loc", []) if part != "body") message = first.get("msg", "invalid request") detail = f"{location}: {message}" if location else message return JSONResponse( status_code=422, content={"error": {"code": "VALIDATION_ERROR", "message": detail}} ) @app.exception_handler(HTTPException) async def http_exception_handler(request: Request, exc: HTTPException) -> JSONResponse: detail = exc.detail if isinstance(detail, dict) and "error" in detail: body = detail else: body = {"error": {"code": "HTTP_ERROR", "message": str(detail)}} return JSONResponse(status_code=exc.status_code, content=body) @app.exception_handler(Exception) async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse: LOGGER.exception("unhandled error on %s: %s", request.url.path, exc) return JSONResponse( status_code=500, content={"error": {"code": "INTERNAL", "message": "internal server error"}} ) # ---------------------------------------------------------------- static web client resolved_web_dir = web_dir or os.path.join(current_dir, "web") if os.path.isdir(resolved_web_dir): app.mount("/", StaticFiles(directory=resolved_web_dir, html=True), name="web") @app.on_event("shutdown") async def _shutdown() -> None: for session in list(sessions.values()): session.close() engine.shutdown() return app def create_fastapi_app( engine: RecognitionEngine, scans: Any = None, settings: Optional[ServerSettings] = None, ) -> FastAPI: """Back-compat entry point; the server entry point is create_app().""" return create_app(engine, settings, scans=scans) # --------------------------------------------------------------------------- # entry point # --------------------------------------------------------------------------- def main() -> None: parser = argparse.ArgumentParser(description="Dorayaki - shelf recognition server") parser.add_argument("--host", type=str, default=None) parser.add_argument("--port", type=int, default=None) parser.add_argument("--device", type=str, default=None, help="cpu | mps | cuda | cuda:0") parser.add_argument("--imgsz", type=int, default=None) parser.add_argument("--weights", type=str, default=None) parser.add_argument("--stock-dir", type=str, default=None) parser.add_argument("--classify-workers", type=int, default=None) parser.add_argument("--live-method", type=str, default=None, choices=list(CLASSIFY_METHODS)) parser.add_argument("--live-conf", type=float, default=None) parser.add_argument("--identity", type=str, default=None, choices=["matcher", "gallery"], help="which mechanism names products (default matcher)") parser.add_argument("--live-inflight", type=int, default=None, help="classifications in flight per session (demo walks want a larger budget)") parser.add_argument("--live-attempts", type=int, default=None, help="classification attempts per track") parser.add_argument("--record-frames", type=str, default=None, help="write processed frames and their results under this directory") parser.add_argument("--record-every", type=int, default=None, help="record one frame in every N (default 3)") parser.add_argument("--demo-walk", action="store_true", help="serve the demo walk's pre-recognised shelf as live tracks") parser.add_argument("--no-tags", action="store_true", help="disable price-tag OCR") parser.add_argument("--max-sessions", type=int, default=None) parser.add_argument("--ssl-certfile", type=str, default=None) parser.add_argument("--ssl-keyfile", type=str, default=None) parser.add_argument("--log-level", type=str, default=None) args = parser.parse_args() logging.basicConfig( level=(args.log_level or server_settings_from_env().log_level or "info").upper(), format="%(asctime)s %(levelname)s %(name)s: %(message)s", ) server = server_settings_from_env() if args.host is not None: server.host = args.host if args.port is not None: server.port = args.port if args.max_sessions is not None: server.max_sessions = args.max_sessions if args.live_method is not None: server.live_method = args.live_method if args.live_conf is not None: server.live_conf = args.live_conf if args.demo_walk: server.demo_walk = True if args.no_tags: server.live_tags = False if args.ssl_certfile is not None: server.ssl_certfile = args.ssl_certfile if args.ssl_keyfile is not None: server.ssl_keyfile = args.ssl_keyfile if args.log_level is not None: server.log_level = args.log_level engine_settings = settings_from_env() if args.record_frames is not None: server.record_frames = args.record_frames if args.record_every is not None: server.record_every = args.record_every if args.identity is not None: engine_settings.identity_backend = args.identity if args.live_inflight is not None: engine_settings.live_inflight = args.live_inflight if args.live_attempts is not None: engine_settings.live_attempts = args.live_attempts if args.device is not None: engine_settings.device = args.device if args.imgsz is not None: engine_settings.imgsz = args.imgsz if args.weights is not None: engine_settings.weights_path = args.weights if args.stock_dir is not None: engine_settings.stock_dir = args.stock_dir if args.classify_workers is not None: engine_settings.classify_workers = args.classify_workers LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...", engine_settings.device or "auto", engine_settings.imgsz) engine = build_default_engine(engine_settings) warm_ms = engine.warm_up() LOGGER.info( "Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d", engine.sku_count(), engine.catalog_version(), warm_ms, engine.classify_workers, ) app = create_app(engine, server) scheme = "https" if server.ssl_certfile else "http" LOGGER.info("Dorayaki recognition server on %s://%s:%d (token %s)", scheme, server.host, server.port, "required" if server.token else "off") if server.host in ("127.0.0.1", "localhost"): LOGGER.info("Open %s://127.0.0.1:%d in a browser for the web client", scheme, server.port) else: LOGGER.info("LAN mode: live camera needs https (secure context); use --ssl-certfile or a tunnel") import uvicorn uvicorn.run( app, host=server.host, port=server.port, log_level=(server.log_level or "info").lower(), access_log=False, # access log would print ?token=... query strings ws="websockets-sansio", ssl_certfile=server.ssl_certfile, ssl_keyfile=server.ssl_keyfile, ) if __name__ == "__main__": main()