""" 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 binascii from contextlib import asynccontextmanager from dataclasses import dataclass import json import logging import os import re import sys import threading import time import uuid from typing import Any, Dict, List, Optional, Tuple # 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 record import FrameRecorder # noqa: E402 from live import MAX_CROP_BYTES, LiveSession, LiveSettings # noqa: E402 from transport import ( # noqa: E402 BodyLimit, ServiceAdvertiser, advertise_host, bearer_matches, declared_length, is_loopback_host, read_bounded, require_token_for_bind, token_matches, ) 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" MAX_CROP_PIXELS = 4_000_000 def _jpeg_dimensions(data: bytes) -> Optional[Tuple[int, int]]: """Read JPEG SOF dimensions without allocating a decoded image.""" if not data.startswith(b"\xff\xd8"): return None offset = 2 while offset + 4 <= len(data): if data[offset] != 0xff: return None while offset < len(data) and data[offset] == 0xff: offset += 1 if offset >= len(data): return None marker = data[offset] offset += 1 if marker in (0xd8, 0x01) or 0xd0 <= marker <= 0xd7: continue if marker in (0xd9, 0xda) or offset + 2 > len(data): return None length = int.from_bytes(data[offset:offset + 2], "big") if length < 2 or offset + length > len(data): return None if marker in (0xc0, 0xc1, 0xc2, 0xc3, 0xc5, 0xc6, 0xc7, 0xc9, 0xca, 0xcb, 0xcd, 0xce, 0xcf): if length < 7: return None return (int.from_bytes(data[offset + 5:offset + 7], "big"), int.from_bytes(data[offset + 3:offset + 5], "big")) offset += length return None # --------------------------------------------------------------------------- # settings # --------------------------------------------------------------------------- @dataclass class ServerSettings: host: str = "127.0.0.1" port: int = 8080 token: Optional[str] = None # Empty means same-origin only. "*" reflects every caller and is not a default. cors_origins: str = "" max_sessions: int = DEFAULT_MAX_SESSIONS live_method: str = "hybrid" live_conf: float = 0.1 live_tags: bool = True # 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, ) def load_local_token(path: str = "config/local.env") -> None: """Fill DORAYAKI_API_TOKEN from an ignored env file. Never prints the value. ponytail: one KEY=VALUE file, no quoting. Upgrade path is a real env loader once a second secret needs the same file. """ if os.environ.get("DORAYAKI_API_TOKEN"): return try: with open(path, "r", encoding="utf-8") as handle: lines = handle.readlines() except OSError: return for line in lines: key, sep, value = line.strip().partition("=") if key == "DORAYAKI_API_TOKEN" and sep and value: os.environ["DORAYAKI_API_TOKEN"] = value return # --------------------------------------------------------------------------- # 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() @asynccontextmanager async def lifespan(_app: FastAPI): yield for session in list(sessions.values()): if session is not None: session.close() engine.shutdown() app = FastAPI( title="Dorayaki Recognition Server", version="1.0", description="Shelf product detection, tracking and classification for capture clients.", lifespan=lifespan, ) # Same-origin is the default. A configured list is echoed back; "*" is not. cors = [o.strip() for o in settings.cors_origins.split(",") if o.strip() and o.strip() != "*"] if cors: app.add_middleware( CORSMiddleware, allow_origins=cors, allow_credentials=False, allow_methods=["GET", "POST"], allow_headers=["Authorization", "Content-Type"], ) scans = scans if scans is not None else build_legacy_scans(engine) # A live session can be recorded so a real scan becomes data to label later. recorder = FrameRecorder( settings.record_frames, every=settings.record_every, identity_backend=engine.settings.identity_backend, ) if settings.record_frames else None samples = {item["id"]: item for item in get_available_shelf_images("add")} # path -> (mtime_ns, size, width, height). Dropped when the file changes. sample_dimensions: Dict[str, Tuple[int, int, int, int]] = {} sample_dimensions_lock = threading.Lock() sessions: Dict[str, LiveSession] = {} sessions_lock = asyncio.Lock() app.state.engine = engine app.state.sessions = sessions async def authorize(authorization: Optional[str] = Header(default=None)) -> None: if not settings.token: return if not bearer_matches(authorization, settings.token): raise _error("UNAUTHORIZED", "missing or invalid bearer token", 401) # ---------------------------------------------------------------- health @app.get("/v1/health") async def health() -> Dict[str, Any]: loop = asyncio.get_running_loop() def _health() -> Dict[str, Any]: gallery = ( engine.identity_gallery() if getattr(engine.settings, "identity_backend", "") == "gallery" else None ) 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": gallery.view_count() if gallery is not None 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)()), }, "sessions_active": sum(session is not None for session in sessions.values()), "uptime_s": round(engine.uptime_s(), 1), } return await loop.run_in_executor(None, _health) # ---------------------------------------------------------------- catalog @app.get("/v1/catalog", dependencies=[Depends(authorize)]) async def catalog() -> Dict[str, Any]: loop = asyncio.get_running_loop() def _catalog() -> Dict[str, Any]: return {"catalog_version": engine.catalog_version(), "products": engine.catalog()} return await loop.run_in_executor(None, _catalog) @app.get("/v1/catalog/{sku_id}/image", dependencies=[Depends(authorize)]) async def catalog_image(sku_id: str) -> Response: loop = asyncio.get_running_loop() path = await loop.run_in_executor(None, 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) loop = asyncio.get_running_loop() product = await loop.run_in_executor(None, 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} def _sample_dimensions(path: str) -> Optional[Tuple[int, int]]: try: stat = os.stat(path) except OSError: with sample_dimensions_lock: sample_dimensions.pop(path, None) return None key = (stat.st_mtime_ns, stat.st_size) with sample_dimensions_lock: cached = sample_dimensions.get(path) if cached is not None and cached[:2] == key: return cached[2], cached[3] import cv2 image = cv2.imread(path) if image is None: with sample_dimensions_lock: sample_dimensions.pop(path, None) return None width, height = int(image.shape[1]), int(image.shape[0]) with sample_dimensions_lock: sample_dimensions[path] = (key[0], key[1], width, height) return width, height def _sample_listing() -> List[Dict[str, Any]]: listing = [] for item in samples.values(): entry = {"id": item["id"], "label": item["label"], "url": f"/v1/samples/{item['id']}"} dimensions = _sample_dimensions(item["path"]) if dimensions is not None: entry["width"], entry["height"] = dimensions listing.append(entry) return listing @app.get("/v1/samples", dependencies=[Depends(authorize)]) async def list_samples() -> Dict[str, Any]: loop = asyncio.get_running_loop() return {"samples": await loop.run_in_executor(None, _sample_listing)} @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 await asyncio.get_running_loop().run_in_executor(None, 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() announced = declared_length(request.headers.get("content-length")) if announced is not None and announced > MAX_UPLOAD_BYTES: raise _error("PAYLOAD_TOO_LARGE", f"image larger than {MAX_UPLOAD_BYTES} bytes", 413) if content_type == "multipart/form-data": form = await request.form(max_part_size=MAX_UPLOAD_BYTES) 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 read_bounded(upload, MAX_UPLOAD_BYTES) 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 read_bounded(request.stream(), MAX_UPLOAD_BYTES) 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) loop = asyncio.get_running_loop() try: image = await loop.run_in_executor(None, engine.decode_image, payload) except UndecodableImage as exc: raise _error("UNDECODABLE_IMAGE", str(exc), 422) try: result = await loop.run_in_executor( None, lambda: 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: # One submit for the whole still-tag operation, on the same # single-worker executor the live sessions read tags on. The default # pool would let a photo's OCR run beside a live session's Vision # work: pyobjc holds the GIL, so the two would interleave and both # would slow down. `_attach_still_tags` calls `run_pass` itself # (reconcile.attach_tags_to_boxes), never submits, so nothing here # queues behind the worker it is running on. result = await loop.run_in_executor( engine.tag_executor, _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. ponytail: every surviving detection pays one Vision read. There is no cheap pre-filter to add, because a card is only distinguishable by its text - its box is pack-shaped, so the aspect/position bounds in tags.py (measured for real tag regions, 1.8-3.8:1) would exclude the boxes cards actually arrive in. Re-measure boxed-card aspect from real photos before narrowing this; a card that slips through is a product name printed on a price tag. """ 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 @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() loop = asyncio.get_running_loop() rows = await loop.run_in_executor( None, lambda: 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) loop = asyncio.get_running_loop() return await loop.run_in_executor( None, lambda: engine.inspect(stored.image, body.box, method=method) ) # ---------------------------------------------------------------- live stream @app.websocket("/v1/stream") async def stream(websocket: WebSocket) -> None: # Query tokens are rejected. A wrong header is rejected before hello. # A missing header is the browser case: hello goes out, then the first # client text must be {"type":"auth","token":"..."} or the socket closes. if websocket.query_params.get("token"): await websocket.close(code=1008) return header = websocket.headers.get("authorization") if settings.token and header and not bearer_matches(header, settings.token): await websocket.close(code=1008) return loop = asyncio.get_running_loop() async with sessions_lock: if len(sessions) >= settings.max_sessions: await websocket.close(code=1013) return session_id = str(uuid.uuid4()) sessions[session_id] = None # reserve before any disk work session = None try: model_version = await loop.run_in_executor(None, engine.model_version) 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, ), ) await loop.run_in_executor(None, session.refresh_catalog) catalog_version = session._catalog_version if settings.live_tags: await loop.run_in_executor(None, session.enable_tags) async with sessions_lock: sessions[session_id] = session except BaseException: if session is not None: session.close() async with sessions_lock: sessions.pop(session_id, None) raise try: await websocket.accept() send_lock = asyncio.Lock() 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": catalog_version, "model": { "version": 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, "max_crop_bytes": MAX_CROP_BYTES, }, } ) except BaseException: session.close() async with sessions_lock: sessions.pop(session_id, None) raise async def processing_loop() -> None: while not session.closed: pending = await session.take_pending_async() 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: # Disk writes must not sit on the event loop between results. await loop.run_in_executor( None, recorder.write, session_id, frame_id, payload, result ) 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) async def handle_capabilities(message: Dict[str, Any]) -> None: """A client that can answer with a full-resolution crop says so here. Until this arrives the session keeps the protocol-v1 behaviour: no result carries a crop request, so a v1 client is never asked. A capability is believed whole or not at all: a message that does not name the lane with a limit the server can honour is refused and the lane stays off, rather than half-applied on the client's say-so. """ if message.get("crop") is not True or "max_crop_bytes" not in message: await send({"type": "error", "code": "BAD_MESSAGE", "message": "crop must be true and state max_crop_bytes"}) return try: capabilities = session.enable_crops(message["max_crop_bytes"]) except ValueError: await send({"type": "error", "code": "BAD_MESSAGE", "message": "max_crop_bytes must be 1..%d" % MAX_CROP_BYTES}) return await send({"type": "capabilities_ack", "capabilities": capabilities}) async def handle_crop(message: Dict[str, Any]) -> None: """Score a crop the client returned for a track we asked about. base64 travels in a text control message, so the binary frame path (4-byte id + JPEG) is untouched and the frame id space is unchanged. """ request_id = message.get("request_id") encoded = message.get("jpeg") if not isinstance(encoded, str) or not encoded: await send({"type": "crop_error", "request_id": request_id, "code": "CROP_BAD_MESSAGE", "message": "jpeg is required"}) return # Base64 expands ~4/3. Refuse on the wire length before decoding it, # the same way the REST frame route does. if len(encoded) > (MAX_CROP_BYTES * 4) // 3 + 8: await send({"type": "crop_error", "request_id": request_id, "code": "CROP_TOO_LARGE", "message": "crop over the limit"}) return try: crop_bytes = base64.b64decode(encoded, validate=True) except (binascii.Error, ValueError): await send({"type": "crop_error", "request_id": request_id, "code": "CROP_BAD_MESSAGE", "message": "jpeg is not valid base64"}) return dimensions = _jpeg_dimensions(crop_bytes) if dimensions is not None and dimensions[0] * dimensions[1] > MAX_CROP_PIXELS: await send({"type": "crop_error", "request_id": request_id, "code": "CROP_TOO_LARGE", "message": "crop dimensions exceed the limit"}) return reply = await loop.run_in_executor( None, session.submit_crop, request_id, message.get("track_id"), message.get("frame_id"), crop_bytes, ) await send(reply) processing = asyncio.create_task(processing_loop()) authed = not settings.token or bearer_matches( websocket.headers.get("authorization"), settings.token ) async def handle_control(message: Dict[str, Any]) -> bool: """Apply one control message. False closes the socket (auth failed).""" nonlocal authed kind = message.get("type") if not authed: presented = message.get("token") if kind == "auth" else None if isinstance(presented, str) and token_matches(presented, settings.token): authed = True return True await websocket.close(code=1008) return False if kind == "auth": return True 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 True 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) elif kind == "capabilities": await handle_capabilities(message) elif kind == "crop": await handle_crop(message) else: await send({"type": "error", "code": "BAD_MESSAGE", "message": f"unknown type '{kind}'"}) return True 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 not authed: await websocket.close(code=1008) break 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 if not await handle_control(parsed): break 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", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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 not isinstance(raw, str): raw = "" if "," in raw: raw = raw.split(",", 1)[1] # Base64 expands ~4/3. Reject the text before decoding it. if len(raw) > (MAX_FRAME_BYTES * 4) // 3 + 8: raise _error("PAYLOAD_TOO_LARGE", f"frame larger than {MAX_FRAME_BYTES} bytes", 413) loop = asyncio.get_running_loop() try: frame_bytes = ( await loop.run_in_executor(None, lambda: base64.b64decode(raw, validate=True)) if raw else b"" ) except (ValueError, base64.binascii.Error): raise _error("BAD_REQUEST", "frame_bytes_b64 is not valid base64", 400) if len(frame_bytes) > MAX_FRAME_BYTES: raise _error("PAYLOAD_TOO_LARGE", f"frame larger than {MAX_FRAME_BYTES} bytes", 413) request = FrameRequest( frame_id=str(payload.get("frame_id") or ""), frame_bytes=frame_bytes, 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 = await loop.run_in_executor(None, 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}", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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", dependencies=[Depends(authorize)]) 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") app.add_middleware(BodyLimit, max_bytes=MAX_UPLOAD_BYTES) if os.path.isdir(resolved_web_dir): app.mount("/", StaticFiles(directory=resolved_web_dir, html=True), name="web") 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") 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("--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.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 load_local_token() if not server.token: server.token = os.environ.get("DORAYAKI_API_TOKEN") or None require_token_for_bind(server.host, server.token) 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) if (engine_settings.identity_backend or "").lower() == "gallery": if engine.identity_gallery() is None: engine.shutdown() raise SystemExit( "gallery identity backend requires local DINOv2 weights; refusing to start" ) 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 is_loopback_host(server.host): 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 advertised = advertise_host(server.host) advertiser = ServiceAdvertiser(advertised, server.port) if advertised else None if advertiser is not None and advertiser.start(): LOGGER.info("Advertising _dorayaki._tcp on %s:%d", advertised, server.port) elif advertised: LOGGER.info("No dns-sd on this host; _dorayaki._tcp was not advertised") try: uvicorn.run( app, host=server.host, port=server.port, log_level=(server.log_level or "info").lower(), access_log=False, # request lines can carry Authorization; keep them off ws="websockets-sansio", ssl_certfile=server.ssl_certfile, ssl_keyfile=server.ssl_keyfile, ) finally: if advertiser is not None: advertiser.close() if __name__ == "__main__": main()