feat: a store's fine-tuned head, and the two-stage gate it needs
The head is now part of a store rather than a script: `finetune_identity.py --stock stock-lopia` trains on that store's enrolled views (27 views over 7 products, 12 augmented copies each) and writes `<stock>/identity_head.pt`, which ProductGallery loads automatically. The first attempt exposed a trap worth recording. A discriminative head has no "none of the above": trained with an angular margin to separate seven products, it maps everything confidently onto those seven, including cartons it has never seen. raw features, threshold 0.55 home shelf 78% accepted foreign shelves 5% accepted trained head, threshold 0.40 home shelf 58% accepted foreign shelves 44% accepted Rejection - the gallery's whole point - disappears. The fix needs no retraining: the two mechanisms fail in opposite directions, so raw features decide *whether* to answer (they generalise) and the head decides *what* to answer (it discriminates). Measured: 47% accepted on the home shelf, 5% on foreign ones. Lower coverage than raw-only, which is the honest trade. The head also needs its own threshold (0.40 against the raw 0.55): a trained head's similarity scale is its own, and reusing the raw line rejected almost everything it said. 175 tests pass. Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
4fd5059ef0
commit
2bf916f2aa
4 changed files with 192 additions and 2 deletions
|
|
@ -34,6 +34,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|||
# 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
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -66,19 +68,50 @@ class ProductGallery:
|
|||
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."""
|
||||
import torch
|
||||
|
||||
if not os.path.exists(self._head_path):
|
||||
return False
|
||||
try:
|
||||
payload = torch.load(self._head_path, map_location="cpu", weights_only=False)
|
||||
except Exception:
|
||||
return False
|
||||
projection = torch.nn.Linear(384, 128)
|
||||
projection.load_state_dict(payload["projection"])
|
||||
projection.eval()
|
||||
self._head = {
|
||||
"projection": projection,
|
||||
"prototypes": torch.nn.functional.normalize(payload["prototypes"].float(), dim=1),
|
||||
"class_ids": list(payload["class_ids"]),
|
||||
}
|
||||
return True
|
||||
|
||||
def _load_model(self) -> Any:
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
|
@ -156,6 +189,7 @@ class ProductGallery:
|
|||
]
|
||||
self._matrix = features
|
||||
self._available = True
|
||||
self.load_head()
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------ querying
|
||||
|
|
@ -171,13 +205,42 @@ class ProductGallery:
|
|||
return features[0]
|
||||
|
||||
def match(self, crop: Any, threshold: Optional[float] = None) -> GalleryMatch:
|
||||
"""Nearest enrolled view, with a rejection when nothing is close enough."""
|
||||
"""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
|
||||
import torch
|
||||
|
||||
limit = self.threshold if threshold is None else float(threshold)
|
||||
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:
|
||||
# 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] = {}
|
||||
|
|
@ -196,6 +259,9 @@ class ProductGallery:
|
|||
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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue