tngo/product-recognition-demo
1
0
Fork 0
forked from vnguyen/product-recognition-demo
Code Pull requests Activity
product-recognition-demo/recognization/code/gallery.py
kiyo9w 3614b23a2c fix: serialize gallery inference so Metal cannot be entered twice 
The server segfaulted a few seconds into every live scan that used the crop
lane. The crash report put the fault in libtorch's Metal shader library
(`MetalShaderLibrary::exec_unary_kernel` inside an unordered_map insert),
reached from `copy_cast_kernel_mps` under `clone`/`reshape`.

Reproduced deterministically: eight threads calling `embed()` on a fresh
process crash it, while the same eight threads after one warm forward are
fine. Metal's shader library is process-global and is not thread-safe while
kernels are still being compiled into it, which is exactly what the first
forward does. The detector already serializes its predict for this reason;
the gallery did not.

Two changes:

* A module-level `_MPS_INFER_LOCK` covers the whole MPS region of a forward -
  model load, `.to(device)`, the forward, and the normalize that follows it.
  Normalize dispatches its own Metal kernels, so leaving it outside the lock
  left the same first-compile race open. Module-level because the state being
  protected is the process's, not one gallery's.
* A cache hit skips the embedding forward, so the model used to be built and
  the kernels compiled by whichever classify worker asked first - up to eight
  racing into one process-wide cache. The load path now warms the backbone
  single-threaded under the existing lock, on the cache-hit branch too. The
  readiness path loads the gallery before the socket accepts, so the warmup
  is free. It is Metal-only and best-effort: a CPU gallery has no shader
  library to compile, and a cache hit is served without a usable backbone by
  contract.

The regression test records overlap rather than timing: a fake backbone that
sleeps and counts concurrent entries, run from four threads, asserting the
count never exceeds one. It fails on the previous gallery.py and needs no MPS.
2026-09-28 17:22:55 +09:00

680 lines
29 KiB
Python
Raw Blame History

"""
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 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)
View git blame Copy permalink