""" 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 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 @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, ) -> None: self.stock_dir = stock_dir self.sku_ids = list(sku_ids) self.image_size = int(image_size) self.threshold = float(threshold) self.device = device self._model = None self._lock = threading.Lock() self._entries: List[GalleryEntry] = [] self._matrix = None self._available = False # ------------------------------------------------------------------ loading 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 = timm.create_model( "vit_small_patch14_dinov2.lvd142m", pretrained=True, num_classes=0, img_size=self.image_size, dynamic_img_size=True, ) 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]: """Every image in the stock directory that belongs to this SKU.""" found: List[str] = [] 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 stem.startswith(f"{sku_id}_") or stem.startswith(f"{sku_id}-"): 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.""" import cv2 import numpy as np import torch paths: List[Tuple[str, str]] = [] for sku_id in self.sku_ids: for path in self._views_for(sku_id): paths.append((sku_id, path)) if not paths: self._available = False return False try: model = self._load_model() except Exception: self._available = False return False batch = [] for _, path in paths: image = cv2.imread(path) if image is None or image.size == 0: continue batch.append(_prepare(image, self.image_size)) if not batch: self._available = False return False tensor = torch.from_numpy(np.stack(batch)).to(self.device) with torch.no_grad(): features = model(tensor) features = torch.nn.functional.normalize(features, dim=1).cpu().numpy() 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 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 view, 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 = self.threshold if threshold is None else float(threshold) embedding = self.embed(crop) 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 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)