""" A product gallery: identity by retrieval, with several views per SKU. The matcher this replaces scores a crop against one template per SKU. Measured on the demo shelf, that is the wrong shape for a store: * one enrolled view identifies 33 % of facings (matcher) / 57 % (embedding); * two views per SKU lift the embedding to 72 %, three to 75 %; * the matcher gets *worse* with more views (33 % -> 5 %), because its raw inlier score is not calibrated to be aggregated across views - and it cannot even express two views of one SKU, since it keys templates by filename. So the gallery is embedding-based. DINOv2 ViT-S/14 over aspect-preserved crops, cosine similarity, several views per SKU, and a rejection threshold so an unknown product comes back as unknown instead of as the nearest stranger. Cost, measured on an Apple M5 (24 crops, batched and one at a time): 224 px CPU 25 ms/crop MPS 2.4 ms/crop 518 px CPU 159 ms/crop MPS 3.6 ms/crop A live frame only needs its *new* tracks embedded (2-4 of them), so ~10 ms a frame on the GPU. That is why this can sit in the live loop at all. """ from __future__ import annotations import hashlib import os import re import threading from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Sequence, Tuple # Below this cosine similarity the nearest view is not the same product. Chosen # from the measured distribution: same-product pairs sit above 0.6, different # products below 0.5, and the gap is where a rejection belongs. DEFAULT_THRESHOLD = 0.55 # The same line for a fine-tuned head, whose similarity scale is its own. DEFAULT_HEAD_THRESHOLD = 0.40 # DINOv2 ViT-S/14, LVD-142M. timm's default source is the Hugging Face copy of # these weights (Apache-2.0); pretrained=True fetches it even when YOLO_OFFLINE=1. # The pin is the sha256 of that model.safetensors (timm hub snapshot 4610ca14). # Not in Git (~84 MB). Place it at weights/dinov2_vits14_lvd142m.safetensors, or # point DORAYAKI_DINOV2_WEIGHTS at a copy. See README "DINOv2 weights". 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 # One thread at a time runs the gallery backbone. Metal's shader library is # process-global and is not thread-safe while kernels are still being compiled # into it, which is what the first forward on a fresh process does; two threads # reaching it together segfault the whole server. Module-level because the state # being protected is the process's, not one gallery's. # ponytail: serializes every embed. Audit C5's batching embedding worker is the # upgrade path, and a second MPS consumer (--device mps puts YOLO on Metal too) # would need to take this same lock rather than its own. _MPS_INFER_LOCK = threading.Lock() def dinov2_weights_path() -> str: override = os.environ.get("DORAYAKI_DINOV2_WEIGHTS") if override: return override here = os.path.dirname(os.path.abspath(__file__)) return os.path.join(os.path.dirname(here), "weights", DINOV2_FILENAME) def _sha256(path: str) -> str: import hashlib digest = hashlib.sha256() with open(path, "rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def load_dinov2(image_size: int) -> Any: """Build the gallery backbone from the pinned local file. Never downloads. A missing or altered file raises FileNotFoundError naming the path and the expected sha256. The head in finetune_identity.py is a Linear(384, 128) on this model's features; a different backbone will not load that head. """ import timm from safetensors.torch import load_file from timm.models.vision_transformer import checkpoint_filter_fn path = dinov2_weights_path() if not os.path.isfile(path): raise FileNotFoundError( f"DINOv2 weights not found at {path}. " f"Copy model.safetensors (sha256 {DINOV2_SHA256}) from the local " f"Hugging Face cache of timm/{DINOV2_MODEL}, or set DORAYAKI_DINOV2_WEIGHTS. " f"The gallery does not download weights." ) digest = _sha256(path) if digest != DINOV2_SHA256: raise FileNotFoundError( f"DINOv2 weights at {path} are sha256 {digest}, expected {DINOV2_SHA256}." ) model = timm.create_model( DINOV2_MODEL, pretrained=False, num_classes=0, img_size=int(image_size), dynamic_img_size=True, ) # The file is the 518px checkpoint. timm's own filter resizes pos_embed when # the caller asks for another size; pretrained=True would do this and also # download, so it stays off. state = checkpoint_filter_fn(load_file(path), model) missing, unexpected = model.load_state_dict(state, strict=False) unexpected = [key for key in unexpected if not key.startswith("head.")] if missing or unexpected: raise RuntimeError( f"DINOv2 weights at {path} do not match {DINOV2_MODEL} " f"(missing={missing}, unexpected={unexpected})." ) return model @dataclass class GalleryEntry: sku_id: str view: str embedding: Any = None @dataclass class GalleryMatch: sku_id: Optional[str] similarity: float runner_up: Optional[str] = None runner_up_similarity: float = 0.0 rejected: bool = False @property def margin(self) -> float: return self.similarity - self.runner_up_similarity def head_class_set(sku_ids: Sequence[str]) -> Optional[List[str]]: """Catalog order, or None when an id is empty or repeated.""" cleaned = [str(sku_id).strip() for sku_id in sku_ids] if not cleaned or any(not sku_id for sku_id in cleaned): return None if len(set(cleaned)) != len(cleaned): return None return cleaned def catalog_digest(stock_dir: str, sku_ids: Sequence[str], csv_name: str = "stock.csv") -> Optional[str]: """Pin for identity_head.pt: stock.csv bytes plus the catalog id list. Image sizes are not included, so another view of the same SKU does not invalidate a head. A changed row or id does. None if the id list is unsafe. """ cleaned = head_class_set(sku_ids) if cleaned is None: return None digest = hashlib.sha256() csv_path = os.path.join(stock_dir, csv_name) try: with open(csv_path, "rb") as handle: digest.update(handle.read()) except OSError: digest.update(b"missing-csv") digest.update(b"\n") digest.update("\n".join(sorted(cleaned)).encode("utf-8")) return digest.hexdigest() class ProductGallery: """Several views per SKU, embedded once, queried by cosine similarity.""" def __init__( self, stock_dir: str, sku_ids: Sequence[str], device: Optional[str] = None, image_size: int = 518, threshold: float = DEFAULT_THRESHOLD, head_threshold: float = DEFAULT_HEAD_THRESHOLD, ) -> None: self.stock_dir = stock_dir self.sku_ids = list(sku_ids) self.image_size = int(image_size) self.threshold = float(threshold) # A trained head produces a different similarity scale than raw features, so # it needs its own line. Measured on a real milk shelf: the head answers 18 of # 72 facings above 0.4 and is right on 94 % of them, while the raw threshold # (0.55) rejects almost everything the head says. self.head_threshold = float(head_threshold) self.device = device self._model = None self._lock = threading.Lock() 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 %. self._head = None self.head_reason = "not_loaded" self._head_path = os.path.join(stock_dir, "identity_head.pt") # ------------------------------------------------------------------ loading def load_head(self) -> bool: """Load the store's fine-tuned head when its labels are this catalog's. weights_only: the file is tensors and strings, not executable code. A missing, malformed, or foreign head leaves the raw gallery in place and records head_reason. Class ids must be unique, nonempty, and a subset of this gallery: a head may cover only the SKUs it was trained on. A catalog_digest, when present, must match stock.csv plus that subset. A head written before the pin loads with head_reason legacy_unpinned; it still cannot name a SKU outside this catalog. """ import torch self._head = None if not os.path.exists(self._head_path): self.head_reason = "missing" return False try: payload = torch.load(self._head_path, map_location="cpu", weights_only=True) reason = self._accept_head(payload, torch) except Exception: self._head = None self.head_reason = "malformed" return False if not reason.startswith("ok"): self._head = None self.head_reason = reason return False self.head_reason = "legacy_unpinned" if reason == "ok:legacy_unpinned" else "ok" return True def _accept_head(self, payload: Any, torch: Any) -> str: if not isinstance(payload, dict): return "malformed" projection_state = payload.get("projection") prototypes = payload.get("prototypes") class_ids = payload.get("class_ids") if not isinstance(projection_state, dict) or not torch.is_tensor(prototypes): return "malformed" if not isinstance(class_ids, (list, tuple)) or prototypes.ndim != 2: return "malformed" if prototypes.shape[0] != len(class_ids) or prototypes.shape[1] != 128: return "malformed" labels = [] for sku_id in class_ids: if not isinstance(sku_id, str) or not sku_id.strip(): return "malformed: identity head class_ids must be nonempty strings" labels.append(sku_id.strip()) if len(set(labels)) != len(labels): return "malformed: identity head class_ids are not unique" expected = head_class_set(self.sku_ids) if expected is None or not set(labels) <= set(expected): return "catalog_mismatch: head class_ids are not a subset of this catalog; retrain finetune_identity.py --stock" recorded = payload.get("catalog_digest", None) if recorded is None: reason = "ok:legacy_unpinned" else: pinned = catalog_digest(self.stock_dir, labels) if pinned is None or recorded != pinned: return "catalog_mismatch: head catalog_digest does not match stock.csv and its class_ids; retrain finetune_identity.py --stock" reason = "ok" if payload.get("base_model") != DINOV2_MODEL: return f"backbone_mismatch: head base_model is not {DINOV2_MODEL}" if type(payload.get("image_size")) is not int or payload.get("image_size") != self.image_size: return f"image_size_mismatch: head image_size is not {self.image_size}" projection = torch.nn.Linear(384, 128) projection.load_state_dict(projection_state) projection.eval() self._head = { "projection": projection, "prototypes": torch.nn.functional.normalize(prototypes.float(), dim=1), "class_ids": labels, } return reason def _load_model(self) -> Any: if self._model is not None: return self._model import torch device = self.device if device is None: device = "mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu" model = load_dinov2(self.image_size) model = model.eval().to(device) self.device = device self._model = model return model def available(self) -> bool: return self._available def _views_for(self, sku_id: str) -> List[str]: """Images whose stem is the SKU or SKU plus _N / -N. A bare prefix would enroll sku_001 as a view of sku_00. """ found: List[str] = [] extra = re.compile(rf"{re.escape(sku_id)}[_-]\d+$") try: names = sorted(os.listdir(self.stock_dir)) except OSError: return found for name in names: stem, extension = os.path.splitext(name) if extension.lower() not in (".jpeg", ".jpg", ".png"): continue if stem == sku_id or extra.fullmatch(stem): found.append(os.path.join(self.stock_dir, name)) return found 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) ] 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. """ 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 _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.floating) or not np.isfinite(matrix).all(): return None # Rows walk the enrolled list forward and never backwards, so a # moved, repeated or foreign row has nowhere left to match. Views # that could not be decoded are simply absent: everyone after one # keeps its order, which is what stops a valid key from renaming a # SKU. expected = [(sku_id, os.path.basename(view)) for sku_id, view in enrolled] rows: List[Tuple[str, str]] = [] position = 0 for sku_id, view in zip(skus, views): while position < len(expected) and expected[position] != (sku_id, view): position += 1 if position == len(expected): return None rows.append((sku_id, os.path.join(self.stock_dir, view))) position += 1 if not self._weights_pinned(): return None return rows, 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 rows: List[Tuple[str, str]] = [] kept: List[Any] = [] 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: return None try: with _MPS_INFER_LOCK: model = self._load_model() tensor = torch.from_numpy(np.stack(kept)).to(self.device) with torch.no_grad(): 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 # A cache hit skips the embedding forward, so without this the model # is built and Metal's kernels are compiled by whichever classify # worker asks first - up to eight racing into one process-wide cache. # The readiness path loads the gallery before the socket accepts, so # the warmup is free and no worker meets a cold model. # # Best-effort on purpose: a cache hit is served without a usable # backbone by contract, and _MPS_INFER_LOCK already makes the lazy # path safe. This only spares the first worker the compile. self._warm_backbone() self._entries = [ GalleryEntry(sku_id=sku_id, view=path, embedding=features[index]) 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 # ------------------------------------------------------------------ querying def embed(self, crop: Any) -> Any: import numpy as np import torch with _MPS_INFER_LOCK: model = self._load_model() tensor = torch.from_numpy(np.stack([_prepare(crop, self.image_size)])).to(self.device) with torch.no_grad(): features = model(tensor) features = torch.nn.functional.normalize(features, dim=1).cpu().numpy() return features[0] def _warm_backbone(self) -> None: """Build the backbone and compile its kernels, single-threaded. Called under self._lock from the load path, so the lock order is self._lock then _MPS_INFER_LOCK - the same order _embed_views takes, and no path takes them reversed. The forward is a dummy canvas: what is being warmed is the model and the shader library, not any one crop. Metal only. A CPU gallery has no shader library to compile and a cache hit is served without touching the backbone by contract, so warming one would spend a model load to protect against nothing. """ import numpy as np if self._resolved_device() != "mps": return try: self.embed(np.zeros((self.image_size, self.image_size, 3), dtype=np.uint8)) except Exception: pass def _resolved_device(self) -> str: """The device the backbone will run on, without building it.""" if self.device is not None: return self.device import torch if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): return "mps" return "cpu" 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: return GalleryMatch(sku_id=None, similarity=0.0, rejected=True) import numpy as np limit = threshold if limit is None: limit = self.head_threshold if self._head is not None else self.threshold limit = float(limit) embedding = self.embed(crop) if self._head is not None: import torch # 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 # well (73 % balanced) but is confidently wrong on strangers (44 % false # accepts), because a discriminative head has no "none of the above". # So the raw features decide *whether* to answer, and the head decides # *what* to answer. raw_similarity = self._matrix @ embedding best_view = float(raw_similarity.max()) if len(raw_similarity) else 0.0 if best_view < self.threshold: return GalleryMatch(sku_id=None, similarity=best_view, rejected=True) with torch.no_grad(): projected = torch.nn.functional.normalize( self._head["projection"](torch.from_numpy(embedding).float()), dim=0 ) similarity = (self._head["prototypes"] @ projected).numpy() ranked = sorted(zip(self._head["class_ids"], similarity), key=lambda kv: kv[1], reverse=True) top_sku, top_score = ranked[0] runner_sku, runner_score = ranked[1] if len(ranked) > 1 else (None, 0.0) if float(top_score) < limit: return GalleryMatch(sku_id=None, similarity=float(top_score), runner_up=top_sku, runner_up_similarity=float(top_score), rejected=True) return GalleryMatch(sku_id=top_sku, similarity=float(top_score), runner_up=runner_sku, runner_up_similarity=float(runner_score)) similarity = self._matrix @ embedding # Several views per SKU: a SKU's score is its best view. order, grouped, starts = self._rank_index() if not order: return GalleryMatch(sku_id=None, similarity=0.0, rejected=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: return GalleryMatch(sku_id=None, similarity=top_score, runner_up=top_sku, runner_up_similarity=top_score, rejected=True) return GalleryMatch(sku_id=top_sku, similarity=top_score, runner_up=runner_sku, runner_up_similarity=runner_score) def rank(self, crop: Any, limit: Optional[int] = None) -> List[Tuple[str, float]]: """Every enrolled SKU by similarity, best first, with no rejection. `match()` answers "is anything close enough"; a shortlist needs "what are the closest candidates", including when none of them is close. Gating a shortlist at the raw threshold empties it for exactly the crops the verifier exists to rescue. """ if not self._available or self._matrix is None: return [] import numpy as np embedding = self.embed(crop) if self._head is not None: import torch with torch.no_grad(): projected = torch.nn.functional.normalize( self._head["projection"](torch.from_numpy(embedding).float()), dim=0 ) similarity = (self._head["prototypes"] @ projected).numpy() ranked = sorted(zip(self._head["class_ids"], similarity.tolist()), key=lambda kv: kv[1], reverse=True) else: similarity = self._matrix @ embedding order, grouped, starts = self._rank_index() if not order: return [] best_by_sku = np.maximum.reduceat(similarity[grouped], starts) ranked = sorted(zip(order, best_by_sku.tolist()), key=lambda item: item[1], reverse=True) return [(str(sku), float(score)) for sku, score in ranked[:limit]] if limit else \ [(str(sku), float(score)) for sku, score in ranked] def head_loaded(self) -> bool: return self._head is not None def view_count(self) -> int: return len(self._entries) def sku_view_counts(self) -> Dict[str, int]: counts: Dict[str, int] = {} for entry in self._entries: counts[entry.sku_id] = counts.get(entry.sku_id, 0) + 1 return counts def _prepare(image: Any, size: int) -> Any: """Aspect-preserving square canvas: the crop's shape is part of the product.""" import cv2 import numpy as np mean = np.array([0.485, 0.456, 0.406], np.float32) std = np.array([0.229, 0.224, 0.225], np.float32) rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) height, width = rgb.shape[:2] side = max(height, width) canvas = np.zeros((side, side, 3), np.uint8) canvas[(side - height) // 2:(side - height) // 2 + height, (side - width) // 2:(side - width) // 2 + width] = rgb resized = cv2.resize(canvas, (size, size), interpolation=cv2.INTER_AREA).astype(np.float32) / 255.0 return ((resized - mean) / std).transpose(2, 0, 1)