385 lines
16 KiB
Python
385 lines
16 KiB
Python
"""
|
|
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)
|