diff --git a/recognization/.gitignore b/recognization/.gitignore index 013e35c..3baa5c9 100644 --- a/recognization/.gitignore +++ b/recognization/.gitignore @@ -22,5 +22,6 @@ runs/detect/train/weights/last.pt weights/dinov2_vits14_lvd142m.safetensors *.log stock-*/identity_head.pt +.gallery_cache_*.npz .demo-token config/local.env diff --git a/recognization/README.md b/recognization/README.md index bcff2d4..20e4f23 100644 --- a/recognization/README.md +++ b/recognization/README.md @@ -64,6 +64,15 @@ cp "$HOME/.cache/huggingface/hub/models--timm--vit_small_patch14_dinov2.lvd142m/ `DORAYAKI_DINOV2_WEIGHTS` overrides that path. A file with any other sha256 is refused. +With the gallery identity backend, enrolled views are embedded into an ignored +`stock-dir/.gallery_cache_.npz` feature cache. The cache is reused +only when its key matches the preprocessing version, image size, pinned DINOv2 +SHA-256, catalog CSV, enrolled SKU/view order and image bytes; its matrix shape, +numeric finite values and row identities are also checked. A stale or malformed +cache is rebuilt. A cache hit still checks the local model file's SHA-256, and +gallery startup refuses readiness if the pinned weights or usable gallery are +missing. OCR continues on the separate tag worker (`tag_executor`). + (`requirements-dev.txt` = `requirements.txt` + `httpx` cho test.) ### Bước 4: Khởi chạy recognition server @@ -100,6 +109,11 @@ Camera trên điện thoại chỉ hoạt động khi trang mở bằng **https* Server có REST + WebSocket đầy đủ, tài liệu chi tiết trong [`docs/API.md`](docs/API.md). Swagger UI: **http://127.0.0.1:8080/docs**. +Live WebSocket protocol `1` also supports an optional full-resolution crop lane: +clients advertise `capabilities`, then answer per-track `crop_request` messages +from their buffered capture. Clients without the capability are unchanged; wire +examples, byte limits, expiry and errors are in [`docs/API.md`](docs/API.md#websocket-v1stream). + ```bash # Nhận diện một ảnh kệ curl -X POST --data-binary @image.jpg -H 'Content-Type: image/jpeg' \ diff --git a/recognization/code/app.py b/recognization/code/app.py index c939471..deadce0 100644 --- a/recognization/code/app.py +++ b/recognization/code/app.py @@ -20,6 +20,7 @@ from __future__ import annotations import argparse import asyncio import base64 +import binascii from contextlib import asynccontextmanager from dataclasses import dataclass import json @@ -73,7 +74,7 @@ from engine import ( # noqa: E402 settings_from_env, ) from record import FrameRecorder # noqa: E402 -from live import LiveSession, LiveSettings # noqa: E402 +from live import MAX_CROP_BYTES, LiveSession, LiveSettings # noqa: E402 from transport import ( # noqa: E402 BodyLimit, ServiceAdvertiser, @@ -473,7 +474,16 @@ def create_app( raise _error("BAD_REQUEST", str(exc), 400) if tags: - result = await loop.run_in_executor(None, _attach_still_tags, result, image) + # 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: @@ -485,6 +495,13 @@ def create_app( 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)(): @@ -733,7 +750,10 @@ def create_app( "imgsz": getattr(engine.settings, "imgsz", 640), }, "config": session.config(), - "limits": {"max_frame_bytes": MAX_FRAME_BYTES}, + "limits": { + "max_frame_bytes": MAX_FRAME_BYTES, + "max_crop_bytes": MAX_CROP_BYTES, + }, } ) except BaseException: @@ -797,6 +817,50 @@ def create_app( 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. + """ + try: + capabilities = session.enable_crops(message.get("max_crop_bytes")) + except (TypeError, ValueError): + await send({"type": "error", "code": "BAD_MESSAGE", + "message": "bad 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 + 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( @@ -835,6 +899,10 @@ def create_app( 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 diff --git a/recognization/code/gallery.py b/recognization/code/gallery.py index 1272e3c..5afcb33 100644 --- a/recognization/code/gallery.py +++ b/recognization/code/gallery.py @@ -47,6 +47,13 @@ DEFAULT_HEAD_THRESHOLD = 0.40 DINOV2_MODEL = "vit_small_patch14_dinov2.lvd142m" DINOV2_SHA256 = "04d27f3400d059fc0cfd7d17dd1909a75bf3ea8fb3eeb48b97cb99e57ee20081" DINOV2_FILENAME = "dinov2_vits14_lvd142m.safetensors" +# What load_dinov2 returns: the same 384 the fine-tuned head is Linear(384, 128) +# on, and the width a cached matrix has to have to belong to this backbone. +DINOV2_EMBED_DIM = 384 + +# Bumped when _prepare() or the cache layout changes. A matrix embedded under an +# older preprocessing pass is stale, not merely old, so it must be re-embedded. +CACHE_VERSION = 1 def dinov2_weights_path() -> str: @@ -192,6 +199,10 @@ class ProductGallery: self._entries: List[GalleryEntry] = [] self._matrix = None self._available = False + # The source fingerprint this matrix was built from, and the per-SKU + # segment bounds derived from the entries it describes. + self._loaded_key = "" + self._rank = None # A fine-tuned projection head, if the store has one (see # finetune_identity.py). Measured on a real milk shelf: raw features put 1 % # balanced accuracy on fine-grained products, a trained head 73 %. @@ -312,35 +323,118 @@ class ProductGallery: found.append(os.path.join(self.stock_dir, name)) return found - def load(self) -> bool: - """Embed every enrolled view. Returns False when there is nothing to load. + def _enrolled(self) -> List[Tuple[str, str]]: + """Every (sku, view) that would be embedded, in embed order.""" + return [ + (sku_id, path) + for sku_id in self.sku_ids + for path in self._views_for(sku_id) + ] - One lock covers the whole load: two callers cannot each embed the catalog. + def _source_key(self, enrolled: Sequence[Tuple[str, str]]) -> str: + """A fingerprint of exactly the inputs an embedding depends on. + + The view bytes, not the name and length: a re-encoded JPEG of the same + size moves the embedding, and a wrong row is worse than a slow start. + Empty when a view cannot be read, which no cache may answer for. """ - with self._lock: - if self._available and self._matrix is not None: - return True - return self._load_locked() + digest = hashlib.sha256() + digest.update(f"v{CACHE_VERSION}|{self.image_size}|{DINOV2_SHA256}\n".encode("utf-8")) + try: + with open(os.path.join(self.stock_dir, "stock.csv"), "rb") as handle: + digest.update(handle.read()) + except OSError: + digest.update(b"missing-csv") + for sku_id, path in enrolled: + digest.update(f"\n{sku_id}\x00{os.path.basename(path)}\x00".encode("utf-8")) + try: + with open(path, "rb") as handle: + digest.update(hashlib.sha256(handle.read()).digest()) + except OSError: + return "" + return digest.hexdigest() - def _load_locked(self) -> bool: + def _cache_path(self) -> str: + return os.path.join(self.stock_dir, f".gallery_cache_{self.image_size}.npz") + + def _weights_pinned(self) -> bool: + """The backbone on disk is still the one this matrix was embedded with. + + A cache hit never opens load_dinov2, so without this a deleted or altered + checkpoint would read as a loaded gallery and only fail on the first + match - after the server had already reported itself ready. + """ + path = dinov2_weights_path() + return os.path.isfile(path) and _sha256(path) == DINOV2_SHA256 + + def _cache_read(self, key: str, enrolled: Sequence[Tuple[str, str]]) -> Optional[Tuple[List[Tuple[str, str]], Any]]: + """The cached rows when they are provably this store's, else None.""" + import numpy as np + + try: + with np.load(self._cache_path(), allow_pickle=False) as payload: + if str(payload["key"].item()) != key: + return None + skus = [str(value) for value in payload["skus"].tolist()] + views = [str(value) for value in payload["views"].tolist()] + matrix = payload["matrix"] + if matrix.ndim != 2 or matrix.shape[1] != DINOV2_EMBED_DIM: + return None + if not matrix.shape[0] or matrix.shape[0] != len(skus) or len(skus) != len(views): + return None + if not np.issubdtype(matrix.dtype, np.number) or not np.isfinite(matrix).all(): + return None + # Undecodable views may be omitted; all cached rows must retain the + # exact enrolled SKU/view order or a valid key could name wrong SKUs. + rows = list(zip(skus, views)) + remaining = iter((sku_id, os.path.basename(view)) for sku_id, view in enrolled) + if not all(any(pair == cached for pair in remaining) for cached in rows): + return None + if not self._weights_pinned(): + return None + paths = [(sku, os.path.join(self.stock_dir, view)) for sku, view in rows] + return paths, matrix.astype(np.float32, copy=False) + except (OSError, ValueError, TypeError, KeyError, IndexError): + return None # absent, malformed, or raced with a checkpoint removal + + def _cache_write(self, key: str, rows: Sequence[Tuple[str, str]], matrix: Any) -> None: + """Replace the cache file in one step. A failed write is not a failed load.""" + import numpy as np + + path = self._cache_path() + tmp = f"{path}.{os.getpid()}.tmp" + try: + with open(tmp, "wb") as handle: + np.savez( + handle, + key=np.array(key), + skus=np.array([sku_id for sku_id, _ in rows]), + views=np.array([os.path.basename(view) for _, view in rows]), + matrix=matrix, + ) + os.replace(tmp, path) + except OSError: # read-only store, no space: embed again next time + try: + os.unlink(tmp) + except OSError: + pass + + def _embed_views(self, enrolled: Sequence[Tuple[str, str]]) -> Optional[Tuple[List[Tuple[str, str]], Any]]: + """Embed every readable view once. None when the backbone is unusable.""" import cv2 import numpy as np import torch - paths: List[Tuple[str, str]] = [] + rows: List[Tuple[str, str]] = [] kept: List[Any] = [] - for sku_id in self.sku_ids: - for path in self._views_for(sku_id): - image = cv2.imread(path) - if image is None or image.size == 0: - continue - paths.append((sku_id, path)) - kept.append(_prepare(image, self.image_size)) + for sku_id, path in enrolled: + image = cv2.imread(path) + if image is None or image.size == 0: + continue + rows.append((sku_id, path)) + kept.append(_prepare(image, self.image_size)) if not kept: - self._available = False - self._entries = [] - self._matrix = None - return False + return None try: model = self._load_model() tensor = torch.from_numpy(np.stack(kept)).to(self.device) @@ -348,16 +442,56 @@ class ProductGallery: features = model(tensor) features = torch.nn.functional.normalize(features, dim=1).cpu().numpy() except Exception: - self._available = False - self._entries = [] - self._matrix = None + return None + return rows, features.astype(np.float32, copy=False) + + def load(self) -> bool: + """Embed every enrolled view. Returns False when there is nothing to load. + + One lock covers the whole load: two callers cannot each embed the catalog. + A loaded gallery is kept only while the store still fingerprints the + same, so a replaced or reordered enrollment re-embeds instead of + answering from a matrix that no longer describes it. + """ + with self._lock: + if self._available and self._matrix is not None: + enrolled = self._enrolled() + if enrolled and self._source_key(enrolled) == self._loaded_key: + return True + return self._load_locked() + + def _unload(self) -> None: + self._available = False + self._entries = [] + self._matrix = None + self._rank = None + self._loaded_key = "" + + def _load_locked(self) -> bool: + enrolled = self._enrolled() + if not enrolled: + self._unload() return False + key = self._source_key(enrolled) + cached = self._cache_read(key, enrolled) if key else None + if cached is None: + embedded = self._embed_views(enrolled) + if embedded is None: + self._unload() + return False + rows, features = embedded + if key: # an unreadable view has no key; never write one we cannot read + self._cache_write(key, rows, features) + else: + rows, features = cached self._entries = [ GalleryEntry(sku_id=sku_id, view=path, embedding=features[index]) - for index, (sku_id, path) in enumerate(paths) + for index, (sku_id, path) in enumerate(rows) ] self._matrix = features self._available = True + self._rank = None + self._loaded_key = key self.load_head() return True @@ -373,6 +507,35 @@ class ProductGallery: features = torch.nn.functional.normalize(features, dim=1).cpu().numpy() return features[0] + def _rank_index(self) -> Tuple[List[str], Any, Any]: + """Segment bounds for the per-SKU maximum, in first-appearance order. + + Entries may arrive in any order, so the grouping is derived from them, not + assumed. Memoised against the list it was built from: replacing the list + rebuilds it, and a SKU list is never longer than the views behind it. + """ + import numpy as np + + if self._rank is not None and self._rank[0] is self._entries: + return self._rank[1], self._rank[2], self._rank[3] + order: List[str] = [] + first_seen: Dict[str, int] = {} + rows: List[int] = [] + for entry in self._entries: + index = first_seen.get(entry.sku_id) + if index is None: + index = len(order) + first_seen[entry.sku_id] = index + order.append(entry.sku_id) + rows.append(index) + column = np.asarray(rows, dtype=np.intp) + # A stable sort groups the rows by SKU; reduceat then takes each segment's + # max, which is a SKU's best view whatever order its views were enrolled in. + grouped = np.argsort(column, kind="stable") + starts = np.searchsorted(column[grouped], np.arange(len(order))) + self._rank = (self._entries, order, grouped, starts) + return order, grouped, starts + def match(self, crop: Any, threshold: Optional[float] = None) -> GalleryMatch: """Nearest enrolled product, with a rejection when nothing is close enough.""" if not self._available or self._matrix is None: @@ -386,7 +549,6 @@ class ProductGallery: embedding = self.embed(crop) if self._head is not None: import torch - if self._head is not None: # Two stages, because the two mechanisms fail in opposite directions. # Measured: raw features reject strangers well (5 % false accepts) but # rank fine-grained products badly (2 % top-1); a trained head ranks them @@ -413,14 +575,11 @@ class ProductGallery: runner_up_similarity=float(runner_score)) similarity = self._matrix @ embedding # Several views per SKU: a SKU's score is its best view. - best_by_sku: Dict[str, float] = {} - for index, entry in enumerate(self._entries): - value = float(similarity[index]) - if value > best_by_sku.get(entry.sku_id, -1.0): - best_by_sku[entry.sku_id] = value - if not best_by_sku: + order, grouped, starts = self._rank_index() + if not order: return GalleryMatch(sku_id=None, similarity=0.0, rejected=True) - ranked = sorted(best_by_sku.items(), key=lambda item: item[1], reverse=True) + best_by_sku = np.maximum.reduceat(similarity[grouped], starts) + ranked = sorted(zip(order, best_by_sku.tolist()), key=lambda item: item[1], reverse=True) top_sku, top_score = ranked[0] runner_sku, runner_score = ranked[1] if len(ranked) > 1 else (None, 0.0) if top_score < limit: diff --git a/recognization/code/live.py b/recognization/code/live.py index 523a786..4675f3b 100644 --- a/recognization/code/live.py +++ b/recognization/code/live.py @@ -58,6 +58,14 @@ STATS_WINDOW = 30 # still overlaps the labelled product; a neighbour-slide crosses it. DRIFT_IOU = 0.5 STATE_UNCONFIRMED = "unconfirmed" +# Which evidence produced the newest observation on a track. A label is only ever +# shown for an accepted match, so this says how that match was measured. +SOURCE_FRAME = "frame" +SOURCE_CROP = "crop" +SOURCE_INHERITED = "inherited" +# A crop the client returns is base64 inside a text control message, so it is +# bounded on its own rather than riding the binary frame path. +MAX_CROP_BYTES = 1024 * 1024 @dataclass @@ -79,6 +87,18 @@ class LiveSettings: glare_fraction: float = 0.02 glare_pixel: int = 245 too_far_height_ratio: float = 0.08 + # Full-resolution crop lane. Off until a client advertises it, so a + # protocol-v1 client is never asked for a crop and its traffic is unchanged. + crops: bool = False + max_crop_bytes: int = 128 * 1024 + # A crop answered this many processed frames late describes a product that + # may have moved on, so it is dropped instead of relabelling the track. + crop_ttl_frames: int = 6 + # Requests outstanding at once. Bounds what one phone can owe us. + max_crop_inflight: int = 4 + # Ask for a native-resolution crop while the box is still this small. At the + # camera's own resolution it may still be large enough to classify. + crop_below_px: int = 64 @dataclass @@ -115,6 +135,11 @@ class TrackState: # or pruned then re-created). Restarting resets the label, so it is always a # re-verification rather than an unexplained label change. restarts: int = 0 + # A crop is asked for at most once per object. A recycled id gets a fresh + # TrackState, so the new object may be asked again. + crop_requested: bool = False + # Which evidence produced the newest observation: frame or crop. + last_source: str = SOURCE_FRAME def area(self) -> int: if not self.raw_box: @@ -137,6 +162,19 @@ def _iou(a: Sequence[int], b: Sequence[int]) -> float: return inter / union if union > 0 else 0.0 +class _CropRejected(Exception): + """A returned crop that may not be used, carrying the reply to send back.""" + + def __init__(self, code: str, message: str): + super().__init__(message) + self.reply = _crop_error(code, message) + + +def _crop_error(code: str, message: str) -> Dict[str, Any]: + """The one shape a refused crop answers with, recoverable like any other.""" + return {"type": "crop_error", "code": code, "message": message} + + class LiveSession: """Stateful live session for one streaming client.""" @@ -195,6 +233,14 @@ class LiveSession: self._frame_origin = [0.0, 0.0] self._shift_map = None self._compliance_tracker = None + # Full-resolution crop lane. A request pins the frame and the track + # object it was minted for, so a late or borrowed crop cannot relabel a + # product the client never photographed. + self._crop_requests: Dict[str, Dict[str, Any]] = {} + self._crop_counter = 0 + self._crops_requested = 0 + self._crops_accepted = 0 + self._crops_rejected = 0 self._last_compliance: Dict[str, Any] = {"violations": 0, "by_rule": {}, "rows": []} self._label_memory: "collections.deque[Dict[str, Any]]" = collections.deque( maxlen=LABEL_MEMORY_MAX @@ -325,6 +371,10 @@ class LiveSession: track.in_flight = 0 track.last_attempt_seq = 0 track.last_attempt_area = 0 + # Outstanding crops were minted against the old method and the + # old generation; answering one now would be evidence measured + # under settings the session no longer holds. + self._crop_requests.clear() return self.config() def config(self) -> Dict[str, Any]: @@ -348,6 +398,7 @@ class LiveSession: self._restart_count.clear() self._slot_history.clear() self._label_memory.clear() + self._crop_requests.clear() self._slots = [] self._latest_image = None self._latest_frame_id = None @@ -492,10 +543,16 @@ class LiveSession: # releasing it. add_done_callback can run the callback inline when the # future already finished, and the callback takes this same lock. submissions = self._plan_classification(image, seq) + crop_requests = self._plan_crop_requests(seq, int(frame_id)) summary = self._summary_locked() stats = self._stats_locked() track_payload = [self._track_payload(state) for state in self._tracks.values()] + if crop_requests: + for row in track_payload: + request = crop_requests.get(row["track_id"]) + if request is not None: + row["crop_request"] = request tag_summary = self._apply_tags(seq, track_payload, tag_bands, shelf_offset, image.shape) if tag_summary is not None: summary["price_tags"] = tag_summary["tags"] @@ -552,6 +609,7 @@ class LiveSession: track.state = STATE_RECOGNIZED track.confirmed = True track.label_box = list(track.box) + track.last_source = SOURCE_INHERITED track.identity = { "sku_id": entry["sku_id"], "confidence": entry["confidence"], @@ -690,6 +748,208 @@ class LiveSession: ) ) + # ------------------------------------------------------------------ crop lane + def capabilities(self) -> Dict[str, Any]: + """What this client is allowed to ask the phone for. + + A client that never sends `capabilities` keeps the protocol-v1 defaults: + no crop lane, so no result ever carries a crop request. + """ + return { + "crop": bool(self.settings.crops), + "max_crop_bytes": int(self.settings.max_crop_bytes), + } + + def enable_crops(self, max_crop_bytes: Optional[int] = None) -> Dict[str, Any]: + """Turn the lane on for a client that advertised it. The client states the + largest crop it will produce; the server keeps the lower of the two.""" + with self._lock: + self.settings.crops = True + if max_crop_bytes is not None: + try: + stated = int(max_crop_bytes) + except (TypeError, ValueError): + stated = 0 + self.settings.max_crop_bytes = max(1024, min(stated, MAX_CROP_BYTES)) + return self.capabilities() + + def _plan_crop_requests(self, seq: int, frame_id: int) -> Dict[int, Dict[str, Any]]: + """Ask the phone for a fuller crop of tracks the frame crop could not use. + + The box is in the pixel space of `frame_id`, the frame the client just + sent, so the phone can answer from the buffer behind that exact frame + rather than from whatever is on screen now. One request per track object, + at most `max_crop_inflight` outstanding, and only for a client that + advertised the lane. + """ + self._expire_crop_requests(seq) + if not self.settings.crops: + return {} + room = self.settings.max_crop_inflight - len(self._crop_requests) + if room <= 0: + return {} + planned: Dict[int, Dict[str, Any]] = {} + for track in self._tracks.values(): + if len(planned) >= room: + break + if track.crop_requested or track.in_flight > 0: + continue + if track.state != STATE_DETECTED or track.confirmed: + continue + if track.hits < self.settings.min_hits or len(track.raw_box) != 4: + continue + x1, y1, x2, y2 = track.raw_box + # The frame crop is too small to say anything about this product. At + # the camera's own resolution it may still be worth a decision, and + # that is the only thing a crop can add here. + if min(x2 - x1, y2 - y1) >= self.settings.crop_below_px: + continue + self._crop_counter += 1 + request_id = "c%d" % self._crop_counter + self._crop_requests[request_id] = { + "request_id": request_id, + "track_id": track.track_id, + "frame_id": int(frame_id), + "object_generation": track.object_generation, + "generation": self._generation, + "seq": seq, + "method": self.settings.method, + } + track.crop_requested = True + # Hold the track while the phone answers, so the frame path does not + # spend a second attempt on the same box at the same time. + track.in_flight += 1 + self._crops_requested += 1 + planned[track.track_id] = { + "request_id": request_id, + "track_id": track.track_id, + "frame_id": int(frame_id), + "box": list(track.raw_box), + } + return planned + + def _expire_crop_requests(self, seq: int) -> None: + """Drop requests the phone did not answer in time. + + A request holds the track busy while it is out, so an unanswered one has + to give that slot back or the track could never be classified at all. + """ + stale = [ + request_id for request_id, record in self._crop_requests.items() + if seq - record["seq"] > self.settings.crop_ttl_frames + ] + for request_id in stale: + record = self._crop_requests.pop(request_id) + self._crops_rejected += 1 + self._release_crop_slot(record) + + def _release_crop_slot(self, record: Dict[str, Any]) -> None: + """Give back the busy slot a crop request took. Lock must be held.""" + track = self._tracks.get(record["track_id"]) + if track is not None and track.object_generation == record["object_generation"]: + track.in_flight = max(0, track.in_flight - 1) + + def submit_crop( + self, + request_id: Any, + track_id: Any, + frame_id: Any, + payload: Optional[bytes], + ) -> Dict[str, Any]: + """Score a crop the client returned for a track we asked about. + + The crop is evidence about one track object, in one frame, decided under + one configuration. Anything that breaks that binding is refused before a + single pixel reaches the classifier, and the frame crop stays the only + evidence for the track. + """ + try: + record, box, payload = self._prepare_crop(request_id, track_id, frame_id, payload) + except _CropRejected as refused: + with self._lock: + self._crops_rejected += 1 + return refused.reply + try: + crop = self.engine.decode_image(payload) + except Exception: + # The busy slot is the request's to give back, whoever failed on it. + with self._lock: + self._release_crop_slot(record) + self._crops_rejected += 1 + return _crop_error("CROP_UNDECODABLE", "crop is not a readable image") + submitted_at = time.perf_counter() + track_id_out = record["track_id"] + generation = record["generation"] + method = record["method"] + object_generation = record["object_generation"] + try: + future = self._executor.submit(self.engine.score_crop, crop, method) + except RuntimeError: # executor already shut down + with self._lock: + self._release_crop_slot(record) + self._crops_rejected += 1 + return _crop_error("CROP_UNAVAILABLE", "classifier is shutting down") + with self._lock: + self._crops_accepted += 1 + future.add_done_callback( + lambda fut, tid=track_id_out, gen=generation, m=method, started=submitted_at, + obj=object_generation, box=box: self._on_classified( + fut, tid, gen, m, started, obj, box, SOURCE_CROP + ) + ) + return { + "type": "crop_ack", + "request_id": record["request_id"], + "track_id": track_id_out, + "frame_id": record["frame_id"], + } + + def _prepare_crop( + self, request_id: Any, track_id: Any, frame_id: Any, payload: Optional[bytes] + ) -> Tuple[Dict[str, Any], List[int], bytes]: + """Check a returned crop against the request it claims to answer. + + Raises [_CropRejected] with the reply to send. The request is consumed + either way: it is single use, so a rejected crop cannot be retried + against a track that has since moved. + """ + with self._lock: + record = self._crop_requests.pop(str(request_id), None) if request_id else None + if record is None: + raise _CropRejected("CROP_UNKNOWN_REQUEST", "no crop is waiting for that request id") + try: + answered_track = int(track_id) + answered_frame = int(frame_id) + except (TypeError, ValueError): + self._release_crop_slot(record) + raise _CropRejected("CROP_BAD_MESSAGE", "track_id and frame_id are required") + if answered_track != record["track_id"] or answered_frame != record["frame_id"]: + # A crop of some other track, or of another frame, says nothing + # about this one however well it matches. + self._release_crop_slot(record) + raise _CropRejected("CROP_MISMATCH", "crop does not match the request it answers") + if self._seq - record["seq"] > self.settings.crop_ttl_frames: + self._release_crop_slot(record) + raise _CropRejected("CROP_STALE", "crop answered too late to describe the track") + if record["generation"] != self._generation: + self._release_crop_slot(record) + raise _CropRejected("CROP_STALE", "the session changed after the request") + track = self._tracks.get(record["track_id"]) + if track is None or track.object_generation != record["object_generation"]: + self._release_crop_slot(record) + raise _CropRejected("CROP_STALE", "the track id was recycled after the request") + # The sender keeps streaming while the requested crop is encoded. + # A newer frame is normal; object generation and the bounded frame + # TTL, not exact latest-sequence equality, guard against reuse. + if track.last_seq < record["seq"]: + self._release_crop_slot(record) + raise _CropRejected("CROP_STALE", "the track is no longer in the requested frame") + limit = self.settings.max_crop_bytes + if not payload or len(payload) > limit: + self._release_crop_slot(record) + raise _CropRejected("CROP_TOO_LARGE", "crop is empty or over %d bytes" % limit) + return record, list(track.box), payload + def _on_classified( self, future: Future, @@ -699,6 +959,7 @@ class LiveSession: submitted_at: float, object_generation: int, submitted_box: List[int], + source: str = SOURCE_FRAME, ) -> None: try: scores = future.result() @@ -715,6 +976,9 @@ class LiveSession: return track.in_flight = max(0, track.in_flight - 1) track.attempts += 1 + # A crop is an attempt against the same budget as a frame crop, so + # the lane cannot be used to re-roll the dice on one product. + track.last_source = source observation = { str(sku_id): float(score) for sku_id, score in (scores or {}).items() } @@ -788,6 +1052,10 @@ class LiveSession: "candidates": list(track.candidates), "attempts": track.attempts, "method": self.settings.method, + # Which evidence produced the newest observation. A client that only + # reads this can tell a label measured on the frame it sent from one + # measured on a crop it sent back. + "source": track.last_source, } product = None if track.state == STATE_RECOGNIZED and track.identity: @@ -865,6 +1133,11 @@ class LiveSession: if self._classify_times: classify_p50 = round(statistics.median(self._classify_times) * 1000.0, 2) return { + "crops": { + "requested": self._crops_requested, + "accepted": self._crops_accepted, + "rejected": self._crops_rejected, + }, "frames_received": self._frames_received, "frames_processed": self._frames_processed, "frames_dropped": self._frames_dropped, diff --git a/recognization/code/ocr.py b/recognization/code/ocr.py index 715d29a..af2823f 100644 --- a/recognization/code/ocr.py +++ b/recognization/code/ocr.py @@ -123,7 +123,8 @@ class VisionOcr(OcrEngine): ok, encoded = cv2.imencode(".png", image) if not ok: return OcrResult(engine=self.name, elapsed_ms=(time.perf_counter() - started) * 1000.0) - data = NSData.dataWithBytes_length_(encoded.tobytes(), len(encoded.tobytes())) + raw = encoded.tobytes() + data = NSData.dataWithBytes_length_(raw, len(raw)) source = self._quartz.CGImageSourceCreateWithData(data, None) cg_image = self._quartz.CGImageSourceCreateImageAtIndex(source, 0, None) diff --git a/recognization/code/test_api.py b/recognization/code/test_api.py index f7cf60a..99dc93d 100644 --- a/recognization/code/test_api.py +++ b/recognization/code/test_api.py @@ -24,6 +24,7 @@ import numpy as np # noqa: E402 from app import ServerSettings, create_app # noqa: E402 from engine import EngineSettings, RecognitionEngine # noqa: E402 +from test_live import SmallBoxDetector # noqa: E402 from test_live import FakeDetector, FakeInventory, FakeMatcher # noqa: E402 from fastapi.testclient import TestClient # noqa: E402 @@ -57,9 +58,9 @@ def make_scans(engine): return ScanService(PerceptionService(adapter)) -def make_app(*, token=None, max_sessions=4): +def make_app(*, token=None, max_sessions=4, detector=None): engine = RecognitionEngine( - FakeDetector(), FakeMatcher(), FakeInventory(), + detector or FakeDetector(), FakeMatcher(), FakeInventory(), settings=EngineSettings(classify_workers=2), ) settings = ServerSettings(token=token, max_sessions=max_sessions) @@ -350,6 +351,109 @@ class WebSocketProtocolTests(unittest.TestCase): missing = socket.receive_json() self.assertEqual(missing["code"], "TRACK_NOT_IN_VIEW") + def crop_app(self): + """A shelf whose only facing is too small for the frame crop to decide.""" + client, engine = make_app(detector=SmallBoxDetector()) + self.addCleanup(engine.shutdown) + return client + + def test_capabilities_turn_on_the_crop_lane(self): + with self.client.websocket_connect("/v1/stream") as socket: + hello = socket.receive_json() + self.assertIn("max_crop_bytes", hello["limits"]) + socket.send_json({"type": "capabilities", "crop": True, "max_crop_bytes": 65536}) + ack = socket.receive_json() + self.assertEqual(ack["type"], "capabilities_ack") + self.assertEqual(ack["capabilities"], {"crop": True, "max_crop_bytes": 65536}) + + def test_a_v1_client_is_never_asked_for_a_crop(self): + with self.client.websocket_connect("/v1/stream") as socket: + socket.receive_json() # hello + for frame_id in (1, 2, 3): + socket.send_bytes(frame_payload(frame_id)) + result = socket.receive_json() + for track in result["tracks"]: + self.assertNotIn("crop_request", track) + self.assertEqual(result["stats"]["crops"]["requested"], 0) + + def test_a_crop_round_trip_labels_the_track(self): + with self.crop_app().websocket_connect("/v1/stream") as socket: + socket.receive_json() # hello + socket.send_json({"type": "capabilities", "crop": True}) + self.assertEqual(socket.receive_json()["type"], "capabilities_ack") + socket.send_bytes(frame_payload(1)) + socket.receive_json() + socket.send_bytes(frame_payload(2)) + result = socket.receive_json() + request = result["tracks"][0].get("crop_request") + self.assertIsNotNone(request, "a facing too small for the frame crop must be asked about") + self.assertEqual(request["frame_id"], 2) + + socket.send_json({ + "type": "crop", + "request_id": request["request_id"], + "track_id": request["track_id"], + "frame_id": request["frame_id"], + "jpeg": base64.b64encode(jpeg_bytes(80, 120)).decode("ascii"), + }) + self.assertEqual(socket.receive_json()["type"], "crop_ack") + socket.send_bytes(frame_payload(3)) + settled = socket.receive_json() + self.assertEqual(settled["stats"]["crops"]["accepted"], 1) + identities = [t["identity"] for t in settled["tracks"]] + self.assertTrue(any(i["source"] == "crop" for i in identities)) + + def test_a_crop_that_does_not_match_its_request_is_refused(self): + with self.crop_app().websocket_connect("/v1/stream") as socket: + socket.receive_json() # hello + socket.send_json({"type": "capabilities", "crop": True}) + socket.receive_json() + socket.send_bytes(frame_payload(1)) + socket.receive_json() + socket.send_bytes(frame_payload(2)) + request = socket.receive_json()["tracks"][0]["crop_request"] + + socket.send_json({ + "type": "crop", + "request_id": request["request_id"], + "track_id": request["track_id"], + "frame_id": 4242, + "jpeg": base64.b64encode(jpeg_bytes(80, 120)).decode("ascii"), + }) + refused = socket.receive_json() + self.assertEqual(refused["type"], "crop_error") + self.assertEqual(refused["code"], "CROP_MISMATCH") + socket.send_bytes(frame_payload(3)) + self.assertEqual(socket.receive_json()["stats"]["crops"]["accepted"], 0) + + def test_a_crop_that_is_not_an_image_is_refused(self): + with self.crop_app().websocket_connect("/v1/stream") as socket: + socket.receive_json() # hello + socket.send_json({"type": "capabilities", "crop": True}) + socket.receive_json() + socket.send_bytes(frame_payload(1)) + socket.receive_json() + socket.send_bytes(frame_payload(2)) + request = socket.receive_json()["tracks"][0]["crop_request"] + + socket.send_json({ + "type": "crop", + "request_id": request["request_id"], + "track_id": request["track_id"], + "frame_id": request["frame_id"], + "jpeg": base64.b64encode(b"not an image").decode("ascii"), + }) + self.assertEqual(socket.receive_json()["code"], "CROP_UNDECODABLE") + + def test_a_crop_without_an_image_is_refused(self): + with self.client.websocket_connect("/v1/stream") as socket: + socket.receive_json() # hello + socket.send_json({"type": "crop", "request_id": "c1", "track_id": 1, "frame_id": 1}) + refused = socket.receive_json() + self.assertEqual(refused["code"], "CROP_BAD_MESSAGE") + socket.send_json({"type": "ping", "t": 1}) + self.assertEqual(socket.receive_json(), {"type": "pong", "t": 1}) + def test_token_is_enforced_on_the_socket(self): """Header or the first post-hello auth frame. A query token is rejected.""" client, engine = make_app(token="s3cret") diff --git a/recognization/code/test_gallery.py b/recognization/code/test_gallery.py index 3c4ddba..eadd11a 100644 --- a/recognization/code/test_gallery.py +++ b/recognization/code/test_gallery.py @@ -32,6 +32,8 @@ from gallery import ( # noqa: E402 load_dinov2, ) +import gallery as gallery_module # noqa: E402 + def a_gallery(views, threshold=0.55): """A gallery with hand-made embeddings: no model, no files.""" @@ -55,6 +57,52 @@ def unit(vector): return (vector / np.linalg.norm(vector)).tolist() +class Backbone: + """A deterministic stand-in for DINOv2: the prepared crop's channel means, + tiled to the width the real backbone returns.""" + + def __init__(self): + self.calls = 0 + + def __call__(self, tensor): + self.calls += 1 + return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3) + + +def a_query(): + """A unit vector in the backbone's space, for match() without a model.""" + query = np.zeros(gallery_module.DINOV2_EMBED_DIM, dtype=np.float32) + query[0] = 1.0 + query[1] = 0.05 + return query / np.linalg.norm(query) + + +class UnreachableBackbone: + def __call__(self, tensor): + raise AssertionError("the gallery loaded a backbone it should have cached") + + +def write_view(path, seed): + import cv2 + + image = np.random.RandomState(seed).randint(0, 256, (24, 32, 3), dtype=np.uint8) + cv2.imwrite(path, image) + + +def a_store(directory, names=("sku_a", "sku_a_1", "sku_b")): + """A store dir: one JPEG per view plus the stock.csv the catalog is read from.""" + with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle: + handle.write("SKU_ID,product_name\n") + for index, name in enumerate(names): + write_view(os.path.join(directory, f"{name}.jpeg"), index) + + +def a_stored_gallery(directory, names, backbone, image_size=32): + gallery = ProductGallery(directory, list(names), device="cpu", image_size=image_size) + gallery._load_model = lambda: backbone + return gallery + + class MatchingTests(unittest.TestCase): def test_the_nearest_view_names_the_product(self): gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))]) @@ -98,6 +146,42 @@ class MatchingTests(unittest.TestCase): close = gallery.match(object()) self.assertGreater(confident.margin, close.margin) + def test_the_per_sku_max_does_not_depend_on_entry_order(self): + # Three SKUs, several views each, enrolled out of order: a SKU is scored + # by its best view, and the answer cannot move because a row moved. + gallery = a_gallery([ + ("sku_b", unit([0, 1, 0])), + ("sku_a", unit([1, 0, 0])), + ("sku_c", unit([0, 0, 1])), + ("sku_b", unit([0.3, 0.95, 0])), + ("sku_a", unit([0.98, 0.2, 0])), + ]) + gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32) + straight = gallery.match(object()) + self.assertEqual(straight.sku_id, "sku_a") + self.assertEqual(straight.runner_up, "sku_b") + + for order in ([4, 0, 3, 1, 2], [2, 3, 1, 4, 0]): + gallery._entries = [gallery._entries[index] for index in order] + gallery._matrix = gallery._matrix[order] + shuffled = gallery.match(object()) + self.assertEqual(shuffled.sku_id, straight.sku_id) + self.assertEqual(shuffled.runner_up, straight.runner_up) + self.assertAlmostEqual(shuffled.similarity, straight.similarity, places=6) + self.assertAlmostEqual(shuffled.runner_up_similarity, straight.runner_up_similarity, places=6) + + def test_a_tie_keeps_the_first_enrolled_sku(self): + # Exactly equidistant: first appearance decides, as it did before the + # per-SKU maximum became a segmented reduction. + gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))]) + gallery.embed = lambda crop: np.array([1.0, 1.0, 0.0], dtype=np.float32) + match = gallery.match(object()) + self.assertFalse(match.rejected) + self.assertEqual(match.sku_id, "sku_a") + self.assertEqual(match.runner_up, "sku_b") + self.assertAlmostEqual(match.similarity, match.runner_up_similarity, places=6) + self.assertAlmostEqual(match.margin, 0.0, places=6) + def test_an_unloaded_gallery_rejects_everything(self): gallery = ProductGallery("does-not-exist", ["sku_a"]) match = gallery.match(object()) @@ -117,6 +201,164 @@ class EnrollmentTests(unittest.TestCase): self.assertEqual(len(gallery._views_for("sku_ab")), 1) +class FeatureCacheTests(unittest.TestCase): + """The cached matrix must be this store's, or it must not be used.""" + + def setUp(self): + import hashlib + import tempfile + + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + a_store(self.directory.name) + # A cache is only served when the pinned backbone is on disk, so point the + # pin at a stand-in file the module's expected digest is told about. The + # check itself - hash this file, compare - is what is under test. + self.weights = os.path.join(self.directory.name, "dinov2.safetensors") + self.write(self.weights, b"not really a checkpoint") + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.weights + self.addCleanup(os.environ.pop, "DORAYAKI_DINOV2_WEIGHTS", None) + self.addCleanup(setattr, gallery_module, "DINOV2_SHA256", gallery_module.DINOV2_SHA256) + gallery_module.DINOV2_SHA256 = hashlib.sha256(b"not really a checkpoint").hexdigest() + + def write(self, path, content): + with open(path, "wb") as handle: + handle.write(content) + return path + + def real_gallery(self): + """A gallery with the real _load_model: the server's own readiness path.""" + return ProductGallery(self.directory.name, ["sku_a", "sku_b"], device="cpu", image_size=32) + + def loaded(self, names=("sku_a", "sku_b"), image_size=32): + backbone = Backbone() + gallery = a_stored_gallery(self.directory.name, names, backbone, image_size=image_size) + self.assertTrue(gallery.load()) + return gallery, backbone + + def test_a_second_start_answers_the_same_from_the_cache(self): + first, backbone = self.loaded() + self.assertEqual(backbone.calls, 1) + self.assertEqual(first.view_count(), 3) + self.assertTrue(os.path.isfile(first._cache_path())) + + cached = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone()) + self.assertTrue(cached.load()) + self.assertEqual(cached.view_count(), first.view_count()) + np.testing.assert_array_equal(cached._matrix, first._matrix) + self.assertEqual([entry.view for entry in cached._entries], + [entry.view for entry in first._entries]) + query = a_query() + first.embed = cached.embed = lambda crop: query + self.assertEqual(vars(cached.match(object())), vars(first.match(object()))) + + def test_changed_view_bytes_re_embed(self): + first, _ = self.loaded() + before = first._matrix.copy() + write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99) + + second, backbone = self.loaded() + self.assertEqual(backbone.calls, 1, "a replaced view must not be served from cache") + self.assertFalse(np.array_equal(second._matrix, before)) + + def test_a_reordered_catalog_does_not_borrow_another_mapping(self): + self.loaded() + swapped, backbone = self.loaded(names=("sku_b", "sku_a")) + self.assertEqual(backbone.calls, 1) + self.assertEqual([entry.sku_id for entry in swapped._entries], ["sku_b", "sku_a", "sku_a"]) + + def test_another_image_size_is_never_served(self): + small, _ = self.loaded(image_size=32) + large, backbone = self.loaded(image_size=64) + self.assertEqual(backbone.calls, 1) + self.assertNotEqual(small._cache_path(), large._cache_path()) + + def test_a_corrupt_cache_re_embeds_the_same_matrix(self): + first, _ = self.loaded() + with open(first._cache_path(), "wb") as handle: + handle.write(b"not an npz at all") + + second, backbone = self.loaded() + self.assertEqual(backbone.calls, 1) + np.testing.assert_array_equal(second._matrix, first._matrix) + + def test_a_tampered_cache_row_is_not_served(self): + # A current key says nothing about the rows inside the file. Swap them + # and the gallery must re-embed, never name a SKU by a moved row. + first, _ = self.loaded() + path = first._cache_path() + with np.load(path, allow_pickle=False) as payload: + key = str(payload["key"].item()) + skus = [str(value) for value in payload["skus"].tolist()] + views = [str(value) for value in payload["views"].tolist()] + matrix = payload["matrix"] + # The two sku_a views are adjacent, so a label swap has to cross SKUs to + # mean anything - and the matrix that stays put is the dangerous one: + # a vector that would then be filed under the wrong product. + self.assertEqual(list(zip(skus, views)), + [("sku_a", "sku_a.jpeg"), ("sku_a", "sku_a_1.jpeg"), ("sku_b", "sku_b.jpeg")]) + tampered = { + "reordered rows": (skus[::-1], views[::-1], matrix[::-1]), + "swapped labels": ([skus[2], skus[1], skus[0]], views, matrix), + "another store's row": (skus[1:] + ["other_store"], views + ["other_store.jpeg"], matrix), + } + for label, (bad_skus, bad_views, bad_matrix) in tampered.items(): + np.savez(path, key=key, skus=np.array(bad_skus), views=np.array(bad_views), matrix=bad_matrix) + backbone = Backbone() + gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone) + self.assertTrue(gallery.load(), label) + self.assertEqual(backbone.calls, 1, f"{label} must not be served from cache") + np.testing.assert_array_equal(gallery._matrix, first._matrix) + self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_a", "sku_b"]) + + def test_a_loaded_gallery_reloads_when_its_source_moves(self): + gallery, backbone = self.loaded() + self.assertTrue(gallery.load()) + self.assertEqual(backbone.calls, 1, "an unchanged store must not re-embed") + + write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99) + self.assertTrue(gallery.load()) + self.assertEqual(backbone.calls, 2) + + with open(os.path.join(self.directory.name, "stock.csv"), "a", encoding="utf-8") as handle: + handle.write("sku_c,new product\n") + self.assertTrue(gallery.load()) + self.assertEqual(backbone.calls, 3, "stock.csv is part of what an embedding is pinned to") + self.assertEqual(gallery.view_count(), 3) + + def test_a_valid_cache_needs_the_pinned_backbone_on_disk(self): + # The regression: a cache hit never opens load_dinov2, so a deleted or + # altered checkpoint used to read as a loaded gallery and then raise on + # the first match - after the server had reported itself ready. + self.loaded() + gallery = self.real_gallery() + self.assertTrue(gallery.load(), "the cache serves while the pin holds") + + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(self.directory.name, "gone.safetensors") + missing = self.real_gallery() + self.assertFalse(missing.load()) + self.assertFalse(missing.available()) + + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.write( + os.path.join(self.directory.name, "altered.safetensors"), b"a different checkpoint" + ) + altered = self.real_gallery() + self.assertFalse(altered.load()) + self.assertFalse(altered.available()) + + def test_an_altered_backbone_re_embeds_rather_than_serves(self): + # With a backbone that can still load, an altered file must re-embed. + self.loaded() + self.assertTrue( + a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone()).load(), + "the cache still serves on the pinned file", + ) + self.write(self.weights, b"a different checkpoint") + backbone = Backbone() + rebuilt = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone) + self.assertTrue(rebuilt.load()) + self.assertEqual(backbone.calls, 1, "an unpinned cache must not be served") + class OfflineWeightsTests(unittest.TestCase): def test_a_missing_weights_file_names_the_path_and_does_not_fetch(self): @@ -246,12 +488,14 @@ class TrainerPinTests(unittest.TestCase): return sentinel - original = finetune_identity.load_dinov2 - finetune_identity.load_dinov2 = fake_load + # build_model resolves the loader from gallery at call time, so that is + # the attribute to replace; finetune_identity has no module-level copy. + original = gallery_module.load_dinov2 + gallery_module.load_dinov2 = fake_load try: model, device = finetune_identity.build_model(224) finally: - finetune_identity.load_dinov2 = original + gallery_module.load_dinov2 = original self.assertIs(model, sentinel) self.assertIn(device, ("cpu", "mps")) self.assertEqual(seen, [224]) diff --git a/recognization/code/test_live.py b/recognization/code/test_live.py index 5b1d02a..dde9d75 100644 --- a/recognization/code/test_live.py +++ b/recognization/code/test_live.py @@ -502,5 +502,179 @@ class ClassificationBudgetTests(unittest.TestCase): engine.shutdown() + +class SmallBoxDetector(FakeDetector): + """A facing too small for the frame crop to decide, so only a crop can.""" + + def __init__(self, box=None): + super().__init__() + self.boxes = [{ + "id": 0, "box": list(box or [100, 100, 130, 200]), "confidence": 0.9, + "class_id": 0, "width": 30, "height": 100, "center": [115, 150], + }] + + +def make_crop_session(matcher=None, settings=None, detector=None): + detector = detector or SmallBoxDetector() + engine = RecognitionEngine( + detector, matcher or FakeMatcher(), FakeInventory(), + settings=__import__("engine").EngineSettings(classify_workers=2), + ) + session = LiveSession("crop", engine, settings=settings or LiveSettings(method="hybrid")) + session.refresh_catalog() + session.enable_crops() + return session, engine, detector + + +def requested_crop(result): + for track in result["tracks"]: + if "crop_request" in track: + return track["crop_request"] + return None + + +class CropLaneTests(unittest.TestCase): + """A full-resolution crop is extra evidence for one track, nothing else.""" + + def test_a_protocol_v1_client_is_never_asked(self): + session, engine, _ = make_crop_session() + session.enable_crops # noqa: B018 - documents that crops start disabled + session.settings.crops = False + process(session, 1) + result = process(session, 2) + self.assertIsNone(requested_crop(result)) + self.assertEqual(result["stats"]["crops"]["requested"], 0) + engine.shutdown() + + def test_accepted_crop_labels_the_track_it_was_asked_about(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + self.assertIsNotNone(request, "a box the frame crop cannot use must be asked about") + self.assertEqual(request["frame_id"], 2) + + reply = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertEqual(reply["type"], "crop_ack") + self.assertTrue(wait_for(lambda: process(session, 3)["stats"]["crops"]["accepted"] == 1)) + engine.shutdown() + + def test_the_label_says_it_was_measured_on_a_crop(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertTrue(wait_for(lambda: "crop_request" not in process(session, 3)["tracks"][0])) + labelled = [t for t in process(session, 4)["tracks"] if t["identity"]["state"] == "recognized"] + self.assertTrue(labelled, "the crop should settle the track the frame could not") + self.assertEqual(labelled[0]["identity"]["sku_id"], "sku_a") + self.assertEqual(labelled[0]["identity"]["source"], "crop") + self.assertEqual(labelled[0]["identity"]["attempts"], 1) + engine.shutdown() + + def test_a_crop_can_arrive_after_the_next_frame(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + process(session, 3) + reply = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertEqual(reply["type"], "crop_ack") + self.assertTrue(wait_for(lambda: session._crops_accepted == 1)) + engine.shutdown() + + def test_a_track_is_asked_once(self): + session, engine, _ = make_crop_session() + process(session, 1) + process(session, 2) + result = process(session, 3) + self.assertIsNone(requested_crop(result), "one crop per track object") + self.assertEqual(result["stats"]["crops"]["requested"], 1) + engine.shutdown() + + def test_a_crop_for_another_frame_is_refused(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + reply = session.submit_crop(request["request_id"], request["track_id"], 999, frame_bytes()) + self.assertEqual(reply["code"], "CROP_MISMATCH") + result = process(session, 3) + self.assertNotEqual(result["tracks"][0]["identity"]["state"], "recognized") + self.assertEqual(result["stats"]["crops"]["rejected"], 1) + self.assertEqual(result["stats"]["crops"]["accepted"], 0) + engine.shutdown() + + def test_a_crop_for_another_track_is_refused(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + reply = session.submit_crop(request["request_id"], request["track_id"] + 41, + request["frame_id"], frame_bytes()) + self.assertEqual(reply["code"], "CROP_MISMATCH") + self.assertEqual(process(session, 3)["stats"]["crops"]["accepted"], 0) + engine.shutdown() + + def test_an_oversized_crop_is_refused(self): + session, engine, _ = make_crop_session(settings=LiveSettings(max_crop_bytes=2048)) + session.enable_crops(2048) + process(session, 1) + request = requested_crop(process(session, 2)) + reply = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], b"\xff\xd8" + b"\x00" * 4096) + self.assertEqual(reply["code"], "CROP_TOO_LARGE") + self.assertEqual(process(session, 3)["stats"]["crops"]["accepted"], 0) + engine.shutdown() + + def test_an_undecodable_crop_is_refused(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + reply = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], b"not an image") + self.assertEqual(reply["code"], "CROP_UNDECODABLE") + self.assertEqual(process(session, 3)["stats"]["crops"]["accepted"], 0) + engine.shutdown() + + def test_a_crop_cannot_be_replayed(self): + session, engine, _ = make_crop_session() + process(session, 1) + request = requested_crop(process(session, 2)) + first = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertEqual(first["type"], "crop_ack") + again = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertEqual(again["code"], "CROP_UNKNOWN_REQUEST") + engine.shutdown() + + def test_an_unknown_request_is_refused(self): + session, engine, _ = make_crop_session() + process(session, 1) + reply = session.submit_crop("c-nope", 1, 1, frame_bytes()) + self.assertEqual(reply["code"], "CROP_UNKNOWN_REQUEST") + engine.shutdown() + + def test_an_unanswered_request_expires_and_frees_the_track(self): + session, engine, _ = make_crop_session(settings=LiveSettings(crop_ttl_frames=1)) + process(session, 1) + request = requested_crop(process(session, 2)) + process(session, 3) + # The request was made on frame 2 with a one-frame budget, so by frame 4 + # it is out of date and the track it held is free again. + result = process(session, 4) + self.assertEqual(result["stats"]["crops"]["rejected"], 1, "an unanswered request is counted") + late = session.submit_crop(request["request_id"], request["track_id"], + request["frame_id"], frame_bytes()) + self.assertEqual(late["code"], "CROP_UNKNOWN_REQUEST") + engine.shutdown() + + def test_the_client_states_the_smaller_crop_limit(self): + session, engine, _ = make_crop_session() + self.assertEqual(session.enable_crops(4096)["max_crop_bytes"], 4096) + self.assertEqual(session.capabilities()["crop"], True) + engine.shutdown() + + if __name__ == "__main__": unittest.main() diff --git a/recognization/docs/API.md b/recognization/docs/API.md index 989c2c0..6a37ee7 100644 --- a/recognization/docs/API.md +++ b/recognization/docs/API.md @@ -224,7 +224,7 @@ Live frames in, tracked products out. One session per connection. {"type": "hello", "protocol": 1, "session_id": "...", "catalog_version": "...", "model": {"version": "...", "device": "cpu", "imgsz": 640}, "config": {"method": "hybrid", "conf": 0.1, "empty_slots": true}, - "limits": {"max_frame_bytes": 4194304}} + "limits": {"max_frame_bytes": 4194304, "max_crop_bytes": 1048576}} ``` **Frames (binary):** `4-byte big-endian frame id` + JPEG bytes. The id is echoed back @@ -236,6 +236,8 @@ so the client can pair a result with the frame it sent. | --- | --- | | `{"type":"config","method":"orb","conf":0.15,"empty_slots":false}` | `{"type":"config_ack","config":{...}}` | | `{"type":"inspect","track_id":12,"method":"hybrid","request_id":"r1"}` | `{"type":"inspect_result","request_id":"r1",...}` or an error | +| `{"type":"capabilities","max_crop_bytes":131072}` | `{"type":"capabilities_ack","capabilities":{"crop":true,"max_crop_bytes":131072}}` | +| `{"type":"crop","request_id":"c1","track_id":12,"frame_id":41,"jpeg":""}` | `{"type":"crop_ack","request_id":"c1","track_id":12,"frame_id":41}` or `crop_error` | | `{"type":"reset"}` | `{"type":"reset_ack"}` | | `{"type":"ping","t":1}` | `{"type":"pong","t":1}` | @@ -257,6 +259,37 @@ so the client can pair a result with the frame it sent. "stats": {"frames_received": 264, "frames_processed": 264, "frames_dropped": 0, "server_fps": 13.2, "classify_pending": 3, "classify_ms_p50": 88.9}} ``` +The optional crop lane is additive to protocol `1`. Without `capabilities`, the +server sends no `crop_request` and existing clients keep their frame/control flow. +After opting in, a `result` may include `"crop_request":{"request_id":"c1", +"track_id":12,"frame_id":41,"box":[812,404,850,460]}` on an eligible +track. The box uses the **sent frame's** pixel coordinates (the raw detection), +not the camera's native resolution. Keep that exact frame's full-resolution +capture briefly, map its box into the capture, encode the crop as JPEG, and reply +in a text JSON `crop` message. Binary frames remain `4-byte id + JPEG`; the crop +does not take a frame id of its own. `crop_ack` means accepted for asynchronous +classification, not recognized; subsequent results carry the updated identity and +`identity.source` (`frame`, `crop`, or `inherited`). `stats.crops` has `requested`, +`accepted`, and `rejected` counts. + +The hello limit is 1 MiB of **decoded crop bytes**. The optional +`max_crop_bytes` in `capabilities` states the client's limit; the acknowledged +session limit is clamped to 1 KiB–1 MiB (default 128 KiB when omitted). Use the +acknowledged limit, not only hello's ceiling. At most four requests are outstanding +per session and each track object is requested at most once. A request expires +after more than six processed frames; clients keep a bounded recent-frame buffer +and may miss a crop if its capture has already been evicted. A late response +cannot label a recycled track. Changing method or resetting invalidates requests. + +Crop failures are recoverable text messages with `type: "crop_error"`, `code`, +and `message`. Pre-decode failures also echo `request_id`; session-level refusals +do not currently echo it. Codes: `CROP_BAD_MESSAGE` (missing/invalid base64 JPEG +or missing IDs), `CROP_TOO_LARGE` (empty or oversized decoded crop, or excessive +encoded length), `CROP_UNKNOWN_REQUEST` (no outstanding request, including one +already expired), `CROP_MISMATCH` (track/frame differs), `CROP_STALE` (late or +invalidated track/session), `CROP_UNDECODABLE` (unreadable image), and +`CROP_UNAVAILABLE` (classifier shutting down). A request is single-use, including +when its response is rejected; continue streaming frames rather than retrying it. **Backpressure:** the server keeps one pending frame per session and processes the newest, counting the ones it replaced in `frames_dropped`. A client should keep one