""" 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 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" 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 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 # 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_path = os.path.join(stock_dir, "identity_head.pt") # ------------------------------------------------------------------ loading def load_head(self) -> bool: """Load the store's fine-tuned head, if it has one. weights_only: the file is tensors and strings, not executable code. A payload that is not that head leaves the raw gallery in place. """ import torch if not os.path.exists(self._head_path): return False try: payload = torch.load(self._head_path, map_location="cpu", weights_only=True) if not isinstance(payload, dict): return False 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 False if not isinstance(class_ids, (list, tuple)) or prototypes.ndim != 2: return False if prototypes.shape[0] != len(class_ids): return False 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": [str(sku_id) for sku_id in class_ids], } except Exception: self._head = None return False return True def _load_model(self) -> Any: if self._model is not None: return self._model import timm 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 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. """ with self._lock: if self._available and self._matrix is not None: return True return self._load_locked() def _load_locked(self) -> bool: import cv2 import numpy as np import torch paths: 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)) if not kept: self._available = False self._entries = [] self._matrix = None return False try: 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: self._available = False self._entries = [] self._matrix = None return False self._entries = [ GalleryEntry(sku_id=sku_id, view=path, embedding=features[index]) for index, (sku_id, path) in enumerate(paths) ] self._matrix = features self._available = True self.load_head() return True # ------------------------------------------------------------------ querying def embed(self, crop: Any) -> Any: import numpy as np import torch 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 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 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 # 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. 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: return GalleryMatch(sku_id=None, similarity=0.0, rejected=True) ranked = sorted(best_by_sku.items(), 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 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)