feat: add bounded crop evidence and keyed gallery cache
This commit is contained in:
kiyo9w
parent
69de9ea418
commit
0d0e9f82b9
8 changed files with 1065 additions and 41 deletions
|
|
@ -47,6 +47,13 @@ DEFAULT_HEAD_THRESHOLD = 0.40
|
|||
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
|
||||
|
||||
|
||||
def dinov2_weights_path() -> str:
|
||||
|
|
@ -192,6 +199,10 @@ class ProductGallery:
|
|||
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 %.
|
||||
|
|
@ -312,35 +323,118 @@ class ProductGallery:
|
|||
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.
|
||||
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)
|
||||
]
|
||||
|
||||
One lock covers the whole load: two callers cannot each embed the catalog.
|
||||
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.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._available and self._matrix is not None:
|
||||
return True
|
||||
return self._load_locked()
|
||||
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 _load_locked(self) -> bool:
|
||||
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.number) or not np.isfinite(matrix).all():
|
||||
return None
|
||||
# Undecodable views may be omitted; all cached rows must retain the
|
||||
# exact enrolled SKU/view order or a valid key could name wrong SKUs.
|
||||
rows = list(zip(skus, views))
|
||||
remaining = iter((sku_id, os.path.basename(view)) for sku_id, view in enrolled)
|
||||
if not all(any(pair == cached for pair in remaining) for cached in rows):
|
||||
return None
|
||||
if not self._weights_pinned():
|
||||
return None
|
||||
paths = [(sku, os.path.join(self.stock_dir, view)) for sku, view in rows]
|
||||
return paths, 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
|
||||
|
||||
paths: List[Tuple[str, str]] = []
|
||||
rows: 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))
|
||||
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:
|
||||
self._available = False
|
||||
self._entries = []
|
||||
self._matrix = None
|
||||
return False
|
||||
return None
|
||||
try:
|
||||
model = self._load_model()
|
||||
tensor = torch.from_numpy(np.stack(kept)).to(self.device)
|
||||
|
|
@ -348,16 +442,56 @@ class ProductGallery:
|
|||
features = model(tensor)
|
||||
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
|
||||
except Exception:
|
||||
self._available = False
|
||||
self._entries = []
|
||||
self._matrix = None
|
||||
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
|
||||
self._entries = [
|
||||
GalleryEntry(sku_id=sku_id, view=path, embedding=features[index])
|
||||
for index, (sku_id, path) in enumerate(paths)
|
||||
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
|
||||
|
||||
|
|
@ -373,6 +507,35 @@ class ProductGallery:
|
|||
features = torch.nn.functional.normalize(features, dim=1).cpu().numpy()
|
||||
return features[0]
|
||||
|
||||
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:
|
||||
|
|
@ -386,7 +549,6 @@ class ProductGallery:
|
|||
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
|
||||
|
|
@ -413,14 +575,11 @@ class ProductGallery:
|
|||
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:
|
||||
order, grouped, starts = self._rank_index()
|
||||
if not order:
|
||||
return GalleryMatch(sku_id=None, similarity=0.0, rejected=True)
|
||||
ranked = sorted(best_by_sku.items(), key=lambda item: item[1], reverse=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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue