892 lines
36 KiB
Python
892 lines
36 KiB
Python
"""
|
|
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 logging
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field, replace
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|
|
|
# Keep the demo box off the network before ultralytics is imported anywhere.
|
|
os.environ.setdefault("YOLO_OFFLINE", "1")
|
|
os.environ.setdefault("YOLO_AUTOINSTALL", "false")
|
|
|
|
COORDINATE_CONVENTION = "pixel-xyxy-origin-top-left"
|
|
REFERENCE_WIDTH = 2048.0
|
|
MAX_UPSCALE = 2.0
|
|
DEFAULT_IMAGE_STORE_MAX = 16
|
|
DEFAULT_IMAGE_STORE_TTL_S = 15 * 60
|
|
DEFAULT_IMAGE_STORE_MAX_BYTES = 256 * 1024 * 1024
|
|
|
|
STATE_DETECTED = "detected"
|
|
STATE_IDENTIFYING = "identifying"
|
|
STATE_RECOGNIZED = "recognized"
|
|
STATE_NEEDS_CONFIRMATION = "needs_confirmation"
|
|
STATE_UNRECOGNIZED = "unrecognized"
|
|
|
|
CLASSIFY_METHODS = ("surf", "orb", "brisk", "hybrid", "hybrid_fast")
|
|
|
|
|
|
LOGGER = logging.getLogger("dorayaki.engine")
|
|
|
|
|
|
class UndecodableImage(ValueError):
|
|
"""Raised when an uploaded payload is not a readable image."""
|
|
|
|
|
|
class UnknownMethod(ValueError):
|
|
"""Raised for a classification method outside CLASSIFY_METHODS."""
|
|
|
|
|
|
class ImageExpired(KeyError):
|
|
"""Raised when an image_id is no longer in the store."""
|
|
|
|
|
|
@dataclass
|
|
class EngineSettings:
|
|
weights_path: str = "runs/detect/train/weights/best.pt"
|
|
stock_dir: str = "stock"
|
|
stock_csv: str = "stock.csv"
|
|
# Measured on this project (Apple M5, bench_engine.py): CPU detects in ~27 ms
|
|
# p50 while MPS takes ~112 ms with identical boxes, so CPU is the default.
|
|
# Override with --device / DORAYAKI_DEVICE (cpu | mps | cuda | cuda:0).
|
|
device: Optional[str] = "cpu"
|
|
# Live-mode defaults, measured with stream_client.py --pace result:
|
|
# hybrid classifies at ~44 ms/crop p50 and is the most accurate method
|
|
# (21/21 SKUs on image.jpg). hybrid_fast (ORB shortlist) is 3.4x faster but
|
|
# agreed with hybrid on only 47% of accepted crops, so it is NOT the default
|
|
# and stays available for experiments only (bench_engine.py --compare).
|
|
# Measured for a moving camera with hybrid classification: 16 in flight keeps
|
|
# the queue draining while detection keeps its share of the CPU. A demo walk
|
|
# over a nearly static shelf wants more; that is a CLI override, not a default.
|
|
live_inflight: int = 16
|
|
live_attempts: int = 2
|
|
# Which mechanism names products: the hand-crafted local-feature matcher, or the
|
|
# embedding gallery with a calibrated rejection. Measured on 54 unseen store
|
|
# photos, the matcher accepts look-alikes (milk cartons as mayonnaise) and its
|
|
# raw inlier score cannot be thresholded; the gallery can say "unknown".
|
|
identity_backend: str = "matcher"
|
|
# Cosine similarity below which the gallery calls a facing unknown. Measured:
|
|
# same-product pairs sit above 0.6, different products below 0.5.
|
|
gallery_threshold: float = 0.55
|
|
gallery_head_threshold: float = 0.40
|
|
gallery_image_size: int = 518
|
|
imgsz: int = 640
|
|
# 0 -> max(2, cpu_count - 2). Measured default below (see build_default_engine).
|
|
classify_workers: int = 5
|
|
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
|
|
# The tag pass for this still (tags, bands, box->tag), kept so a later
|
|
# classify can re-reconcile against the identities it just produced.
|
|
tag_bundle: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
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 set_tag_bundle(self, image_id: str, bundle: Dict[str, Any]) -> None:
|
|
with self._lock:
|
|
stored = self._items.get(image_id)
|
|
if stored is not None:
|
|
stored.tag_bundle = bundle
|
|
|
|
def get_tag_bundle(self, image_id: str) -> Optional[Dict[str, Any]]:
|
|
with self._lock:
|
|
stored = self._items.get(image_id)
|
|
return stored.tag_bundle if stored is not None else None
|
|
|
|
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()
|
|
# The embedding gallery is built on first use (it loads a model and every
|
|
# enrolled view), and cached as False when it cannot be built at all.
|
|
self._gallery: Any = None
|
|
self._gallery_lock = threading.Lock()
|
|
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()
|
|
# OCR is the slowest stage (0.2-3.4 s). It gets its own executor so tag
|
|
# passes can never occupy a classification worker, and a lock so several
|
|
# live sessions do not pile OCR requests onto the CPU at once.
|
|
self.ocr_engine: Optional[Any] = None
|
|
self.ocr_lock = threading.Lock()
|
|
self._tag_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="tags")
|
|
self._owns_tag_executor = True
|
|
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
|
|
|
|
@property
|
|
def tag_executor(self) -> ThreadPoolExecutor:
|
|
return self._tag_executor
|
|
|
|
def set_ocr_engine(self, engine: Any) -> None:
|
|
"""Attach an OCR backend; None disables tag reading."""
|
|
self.ocr_engine = engine
|
|
|
|
def read_tags(self, image: Any) -> Any:
|
|
"""OCR one crop under the shared lock (Vision is happy with concurrency,
|
|
but the CPU is not: one OCR at a time keeps detection responsive)."""
|
|
if self.ocr_engine is None:
|
|
return None
|
|
with self.ocr_lock:
|
|
return self.ocr_engine.read(image)
|
|
|
|
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)
|
|
if self._owns_tag_executor:
|
|
self._tag_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
|
|
@property
|
|
def identity_scale(self) -> str:
|
|
"""What the scores from `score_crop` mean: inlier counts or cosine similarity."""
|
|
return "cosine" if (self.settings.identity_backend or "matcher").lower() == "gallery" else "inlier"
|
|
|
|
def score_crop(self, crop: Any, method: Optional[str] = None) -> Dict[str, float]:
|
|
"""Score one crop. Gallery never falls back to the matcher.
|
|
|
|
A missing gallery returns no scores (unknown), not inlier counts. Mixing
|
|
those scales is how a live scan accepted a look-alike the still path rejected.
|
|
"""
|
|
if self.identity_scale == "cosine":
|
|
gallery = self.identity_gallery()
|
|
if not gallery:
|
|
return {}
|
|
match = gallery.match(crop)
|
|
if match.rejected or not match.sku_id:
|
|
return {}
|
|
scores = {match.sku_id: match.similarity}
|
|
if match.runner_up:
|
|
scores[match.runner_up] = match.runner_up_similarity
|
|
return {str(sku): float(score) for sku, score in scores.items()}
|
|
method = (method or self.settings.live_method).lower()
|
|
if method not in CLASSIFY_METHODS:
|
|
raise UnknownMethod(method)
|
|
if method == "hybrid_fast":
|
|
scores, _ = self.matcher.score_crop_fast(crop)
|
|
else:
|
|
scores, _ = self.matcher.score_crop(crop, method=method)
|
|
return {str(sku): float(score) for sku, score in scores.items()}
|
|
|
|
def identity_gallery(self) -> Any:
|
|
"""The embedding gallery, loaded once (None when unavailable)."""
|
|
if self._gallery is False:
|
|
return None
|
|
if self._gallery is not None:
|
|
return self._gallery
|
|
with self._gallery_lock:
|
|
if self._gallery is False:
|
|
return None
|
|
if self._gallery is not None:
|
|
return self._gallery
|
|
from gallery import ProductGallery
|
|
|
|
gallery = ProductGallery(
|
|
self.settings.stock_dir,
|
|
[product["sku_id"] for product in self.catalog()],
|
|
threshold=self.settings.gallery_threshold,
|
|
head_threshold=self.settings.gallery_head_threshold,
|
|
image_size=self.settings.gallery_image_size,
|
|
)
|
|
try:
|
|
loaded = gallery.load()
|
|
except Exception as exc: # torch/timm missing, no weights, no network
|
|
LOGGER.warning("identity gallery unavailable: %s", exc)
|
|
loaded = False
|
|
# False means tried and failed. Do not retry and do not use the matcher.
|
|
self._gallery = gallery if loaded else False
|
|
if loaded:
|
|
LOGGER.info(
|
|
"identity gallery: %d views over %d SKUs on %s",
|
|
gallery.view_count(),
|
|
len(gallery.sku_ids),
|
|
gallery.device,
|
|
)
|
|
else:
|
|
LOGGER.warning("identity gallery unavailable; gallery requests stay unrecognized")
|
|
return gallery if loaded else None
|
|
|
|
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 []
|
|
|
|
def area_of(box: Dict[str, Any]) -> int:
|
|
if "width" in box and "height" in box:
|
|
return int(box["width"]) * int(box["height"])
|
|
x1, y1, x2, y2 = [int(v) for v in box["box"]]
|
|
return max(0, x2 - x1) * max(0, y2 - y1)
|
|
|
|
ordered = sorted(boxes, key=area_of, reverse=True)
|
|
crops = [self.crop(image, box["box"]) for box in ordered]
|
|
|
|
if (self.settings.identity_backend or "matcher").lower() == "gallery":
|
|
gallery = self.identity_gallery()
|
|
# No matcher fallback: an unavailable gallery is unrecognized, not a different model.
|
|
return self._classify_with_gallery(gallery, ordered, crops)
|
|
if method == "hybrid_fast":
|
|
futures = [self._executor.submit(self.matcher.score_crop_fast, crop) for crop in crops]
|
|
else:
|
|
futures = [self._executor.submit(self.matcher.score_crop, crop, method) for crop in crops]
|
|
results: List[Dict[str, Any]] = []
|
|
for box, future in zip(ordered, futures):
|
|
try:
|
|
scores, details = future.result()
|
|
except Exception as exc: # a single bad crop must not fail the batch
|
|
scores, details = {}, {"error": str(exc)}
|
|
decision = decide_identity(scores, details=details, method=method)
|
|
results.append(
|
|
{
|
|
"detection_id": int(box["id"]),
|
|
"box": list(box["box"]),
|
|
"identity": identity_payload(decision),
|
|
}
|
|
)
|
|
return results
|
|
|
|
def inspect(self, image: Any, box: Sequence[int], method: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Full single-crop result, using the same backend as recognize/live."""
|
|
from identity import decide_identity, decision_for_match
|
|
|
|
crop = self.crop(image, box)
|
|
if self.identity_scale == "cosine":
|
|
gallery = self.identity_gallery()
|
|
match = gallery.match(crop) if gallery else None
|
|
# match() already applied the gallery's own threshold. Ranking again
|
|
# with the published cosine gates must not invent a second line.
|
|
decision = decision_for_match(match)
|
|
return self._inspect_payload(box, "gallery", decision, crop, visualize=False)
|
|
|
|
method = (method or self.settings.live_method).lower()
|
|
if method not in CLASSIFY_METHODS:
|
|
raise UnknownMethod(method)
|
|
if method == "hybrid_fast":
|
|
scores, details = self.matcher.score_crop_fast(crop)
|
|
visual_method = "hybrid"
|
|
else:
|
|
scores, details = self.matcher.score_crop(crop, method=method)
|
|
visual_method = method
|
|
decision = decide_identity(scores, details=details, method=method)
|
|
return self._inspect_payload(box, method, decision, crop, visual_method=visual_method)
|
|
|
|
def _inspect_payload(
|
|
self,
|
|
box: Sequence[int],
|
|
method: str,
|
|
decision: Dict[str, Any],
|
|
crop: Any,
|
|
*,
|
|
visual_method: Optional[str] = None,
|
|
visualize: bool = True,
|
|
) -> Dict[str, Any]:
|
|
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 = {key: dict(value) for key, value 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 visualize and visualized_sku and visual_method:
|
|
try:
|
|
visualization = self.matcher.generate_match_visualization(
|
|
crop, visualized_sku, method=visual_method
|
|
)
|
|
except Exception:
|
|
visualization = None
|
|
|
|
return {
|
|
"box": [int(value) for value 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),
|
|
},
|
|
}
|
|
|
|
def _classify_with_gallery(
|
|
self,
|
|
gallery: Any,
|
|
ordered: Sequence[Dict[str, Any]],
|
|
crops: Sequence[Any],
|
|
) -> List[Dict[str, Any]]:
|
|
"""Identity by retrieval: nearest enrolled view, or an honest unknown.
|
|
|
|
The gallery scores a SKU by its best view and rejects when nothing is close
|
|
enough, which is the behaviour the matcher cannot provide - measured on 54
|
|
unseen store photos, the matcher labelled milk cartons as mayonnaise and its
|
|
score could not be thresholded to stop it.
|
|
"""
|
|
from identity import decision_for_match
|
|
|
|
results: List[Dict[str, Any]] = []
|
|
for box, crop in zip(ordered, crops):
|
|
match = None
|
|
if gallery is not None:
|
|
try:
|
|
match = gallery.match(crop)
|
|
except Exception as exc: # a bad crop must not fail the batch
|
|
match = None
|
|
LOGGER.debug("gallery match failed: %s", exc)
|
|
decision = decision_for_match(match)
|
|
results.append({
|
|
"detection_id": int(box["id"]),
|
|
"box": list(box["box"]),
|
|
"identity": identity_payload(decision),
|
|
})
|
|
return results
|
|
|
|
@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()
|
|
# Detection and classification share the CPU. torch defaults to every core,
|
|
# which starves the classification pool (and the pool then slows detection
|
|
# back). Measured on the M5: capping torch at 4 threads and classifying with
|
|
# 5 workers gave 12.2 results/s single-client (46 ms detect) versus 11.1/s
|
|
# (49 ms) uncapped, and 8.5/s each for two clients. Override with
|
|
# DORAYAKI_TORCH_THREADS (0 = leave torch's default).
|
|
torch_threads = int(os.environ.get("DORAYAKI_TORCH_THREADS", "4") or 0)
|
|
if torch_threads > 0:
|
|
import torch
|
|
|
|
torch.set_num_threads(torch_threads)
|
|
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)
|
|
engine = RecognitionEngine(detector, matcher, inventory, settings=settings)
|
|
|
|
from ocr import build_ocr_engine
|
|
|
|
try:
|
|
ocr_engine = build_ocr_engine()
|
|
except Exception as exc: # a broken OCR install must not stop recognition
|
|
print(f"[Engine] OCR disabled: {exc}")
|
|
ocr_engine = None
|
|
engine.set_ocr_engine(ocr_engine)
|
|
print(f"[Engine] OCR engine: {getattr(ocr_engine, 'name', 'none')}")
|
|
return engine
|
|
|
|
|
|
def settings_from_env(env: Optional[Dict[str, str]] = None) -> EngineSettings:
|
|
"""Build EngineSettings from DORAYAKI_* environment variables."""
|
|
env = dict(os.environ if env is None else env)
|
|
base = EngineSettings()
|
|
return replace(
|
|
base,
|
|
weights_path=env.get("DORAYAKI_WEIGHTS", base.weights_path),
|
|
stock_dir=env.get("DORAYAKI_STOCK_DIR", base.stock_dir),
|
|
device=env.get("DORAYAKI_DEVICE") or base.device,
|
|
imgsz=int(env.get("DORAYAKI_IMGSZ", base.imgsz)),
|
|
classify_workers=int(env.get("DORAYAKI_CLASSIFY_WORKERS", base.classify_workers)),
|
|
live_method=env.get("DORAYAKI_LIVE_METHOD", base.live_method),
|
|
live_conf=float(env.get("DORAYAKI_LIVE_CONF", base.live_conf)),
|
|
live_inflight=int(env.get("DORAYAKI_LIVE_INFLIGHT", base.live_inflight)),
|
|
live_attempts=int(env.get("DORAYAKI_LIVE_ATTEMPTS", base.live_attempts)),
|
|
identity_backend=env.get("DORAYAKI_IDENTITY", base.identity_backend),
|
|
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()
|