product-recognition-demo/recognization/code/gallery.py

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