feat: add bounded crop evidence and keyed gallery cache

This commit is contained in:
kiyo9w 2026-09-28 15:49:34 +09:00
parent 69de9ea418
commit 0d0e9f82b9
8 changed files with 1065 additions and 41 deletions

View file

@ -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

View file

@ -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

View file

@ -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):
for sku_id, path in enrolled:
image = cv2.imread(path)
if image is None or image.size == 0:
continue
paths.append((sku_id, path))
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:
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:

View file

@ -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,

View file

@ -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)

View file

@ -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")

View file

@ -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])

View file

@ -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()