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
|
|
@ -93,6 +93,7 @@ class EngineSettings:
|
||||||
# Cosine similarity below which the gallery calls a facing unknown. Measured:
|
# Cosine similarity below which the gallery calls a facing unknown. Measured:
|
||||||
# same-product pairs sit above 0.6, different products below 0.5.
|
# same-product pairs sit above 0.6, different products below 0.5.
|
||||||
gallery_threshold: float = 0.55
|
gallery_threshold: float = 0.55
|
||||||
|
gallery_head_threshold: float = 0.40
|
||||||
gallery_image_size: int = 518
|
gallery_image_size: int = 518
|
||||||
imgsz: int = 640
|
imgsz: int = 640
|
||||||
# 0 -> max(2, cpu_count - 2). Measured default below (see build_default_engine).
|
# 0 -> max(2, cpu_count - 2). Measured default below (see build_default_engine).
|
||||||
|
|
@ -430,6 +431,7 @@ class RecognitionEngine:
|
||||||
self.settings.stock_dir,
|
self.settings.stock_dir,
|
||||||
[product["sku_id"] for product in self.catalog()],
|
[product["sku_id"] for product in self.catalog()],
|
||||||
threshold=self.settings.gallery_threshold,
|
threshold=self.settings.gallery_threshold,
|
||||||
|
head_threshold=self.settings.gallery_head_threshold,
|
||||||
image_size=self.settings.gallery_image_size,
|
image_size=self.settings.gallery_image_size,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
|
|
|
||||||
|
|
@ -133,13 +133,84 @@ def train_head(features, labels, class_ids, epochs, device, margin: float = 0.0,
|
||||||
return projection, torch.nn.functional.normalize(prototypes.detach(), dim=1)
|
return projection, torch.nn.functional.normalize(prototypes.detach(), dim=1)
|
||||||
|
|
||||||
|
|
||||||
|
def train_from_stock(args) -> None:
|
||||||
|
"""Train the head on a store's own enrolled views, for the gallery to load.
|
||||||
|
|
||||||
|
The head belongs to the store: the views in its stock directory are its products,
|
||||||
|
and a head trained on them is what turns the gallery from a lookup that cannot
|
||||||
|
tell two milk cartons apart into one that can. One to four views per SKU is
|
||||||
|
enough, with augmentation standing in for the rest.
|
||||||
|
"""
|
||||||
|
import glob
|
||||||
|
import json
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from gallery import ProductGallery
|
||||||
|
|
||||||
|
stock = args.stock
|
||||||
|
gallery = ProductGallery(stock, [])
|
||||||
|
# the SKUs are whatever the stock directory holds
|
||||||
|
skus = sorted({os.path.splitext(name)[0].rsplit("_", 1)[0] if os.path.splitext(name)[0].rsplit("_", 1)[-1].isdigit()
|
||||||
|
else os.path.splitext(name)[0]
|
||||||
|
for name in os.listdir(stock)
|
||||||
|
if os.path.splitext(name)[1].lower() in (".jpeg", ".jpg", ".png")})
|
||||||
|
views = {sku: gallery._views_for(sku) for sku in skus}
|
||||||
|
views = {sku: paths for sku, paths in views.items() if paths}
|
||||||
|
if len(views) < 2:
|
||||||
|
print(f"{stock}: needs at least two products with views, found {len(views)}")
|
||||||
|
return
|
||||||
|
|
||||||
|
crops, labels = [], []
|
||||||
|
for sku, paths in sorted(views.items()):
|
||||||
|
for path in paths:
|
||||||
|
image = cv2.imread(path)
|
||||||
|
if image is not None and image.size:
|
||||||
|
crops.append(image)
|
||||||
|
labels.append(sku)
|
||||||
|
class_ids = sorted(views)
|
||||||
|
print(f"{stock}: {len(crops)} views over {len(class_ids)} products")
|
||||||
|
|
||||||
|
model, device = build_model()
|
||||||
|
rng = np.random.RandomState(11)
|
||||||
|
augmented, augmented_labels = [], []
|
||||||
|
for crop, label in zip(crops, labels):
|
||||||
|
augmented.append(crop)
|
||||||
|
augmented_labels.append(label)
|
||||||
|
for _ in range(args.augment):
|
||||||
|
augmented.append(augment(crop, rng))
|
||||||
|
augmented_labels.append(label)
|
||||||
|
print(f"embedding {len(augmented)} crops on {device}...")
|
||||||
|
features = embed(model, device, augmented)
|
||||||
|
|
||||||
|
projection, prototypes = train_head(features, augmented_labels, class_ids,
|
||||||
|
args.epochs, device, margin=0.25)
|
||||||
|
payload = {
|
||||||
|
"class_ids": class_ids,
|
||||||
|
"image_size": 518,
|
||||||
|
"base_model": "vit_small_patch14_dinov2.lvd142m",
|
||||||
|
"projection": {key: value.cpu() for key, value in projection.state_dict().items()},
|
||||||
|
"prototypes": prototypes.cpu(),
|
||||||
|
"views": {sku: len(paths) for sku, paths in sorted(views.items())},
|
||||||
|
}
|
||||||
|
path = os.path.join(stock, "identity_head.pt")
|
||||||
|
torch.save(payload, path)
|
||||||
|
print(f"wrote {path} ({len(class_ids)} products, {len(crops)} views)")
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument("--epochs", type=int, default=60)
|
parser.add_argument("--epochs", type=int, default=60)
|
||||||
parser.add_argument("--views", type=int, default=2, help="enrolment views per class")
|
parser.add_argument("--views", type=int, default=2, help="enrolment views per class")
|
||||||
parser.add_argument("--augment", type=int, default=12, help="augmented copies per enrolment view")
|
parser.add_argument("--augment", type=int, default=12, help="augmented copies per enrolment view")
|
||||||
parser.add_argument("--image", default=os.path.expanduser("~/dorayaki-data/shelf-images"))
|
parser.add_argument("--image", default=os.path.expanduser("~/dorayaki-data/shelf-images"))
|
||||||
|
parser.add_argument("--stock", default=None,
|
||||||
|
help="train on a store's enrolled views instead of the milk shelf, "
|
||||||
|
"and write <stock>/identity_head.pt for the gallery to load")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
if args.stock:
|
||||||
|
train_from_stock(args)
|
||||||
|
return
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
# from the measured distribution: same-product pairs sit above 0.6, different
|
||||||
# products below 0.5, and the gap is where a rejection belongs.
|
# products below 0.5, and the gap is where a rejection belongs.
|
||||||
DEFAULT_THRESHOLD = 0.55
|
DEFAULT_THRESHOLD = 0.55
|
||||||
|
# The same line for a fine-tuned head, whose similarity scale is its own.
|
||||||
|
DEFAULT_HEAD_THRESHOLD = 0.40
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
@ -66,19 +68,50 @@ class ProductGallery:
|
||||||
device: Optional[str] = None,
|
device: Optional[str] = None,
|
||||||
image_size: int = 518,
|
image_size: int = 518,
|
||||||
threshold: float = DEFAULT_THRESHOLD,
|
threshold: float = DEFAULT_THRESHOLD,
|
||||||
|
head_threshold: float = DEFAULT_HEAD_THRESHOLD,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.stock_dir = stock_dir
|
self.stock_dir = stock_dir
|
||||||
self.sku_ids = list(sku_ids)
|
self.sku_ids = list(sku_ids)
|
||||||
self.image_size = int(image_size)
|
self.image_size = int(image_size)
|
||||||
self.threshold = float(threshold)
|
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.device = device
|
||||||
self._model = None
|
self._model = None
|
||||||
self._lock = threading.Lock()
|
self._lock = threading.Lock()
|
||||||
self._entries: List[GalleryEntry] = []
|
self._entries: List[GalleryEntry] = []
|
||||||
self._matrix = None
|
self._matrix = None
|
||||||
self._available = False
|
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
|
# ------------------------------------------------------------------ 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:
|
def _load_model(self) -> Any:
|
||||||
if self._model is not None:
|
if self._model is not None:
|
||||||
return self._model
|
return self._model
|
||||||
|
|
@ -156,6 +189,7 @@ class ProductGallery:
|
||||||
]
|
]
|
||||||
self._matrix = features
|
self._matrix = features
|
||||||
self._available = True
|
self._available = True
|
||||||
|
self.load_head()
|
||||||
return True
|
return True
|
||||||
|
|
||||||
# ------------------------------------------------------------------ querying
|
# ------------------------------------------------------------------ querying
|
||||||
|
|
@ -171,13 +205,42 @@ class ProductGallery:
|
||||||
return features[0]
|
return features[0]
|
||||||
|
|
||||||
def match(self, crop: Any, threshold: Optional[float] = None) -> GalleryMatch:
|
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:
|
if not self._available or self._matrix is None:
|
||||||
return GalleryMatch(sku_id=None, similarity=0.0, rejected=True)
|
return GalleryMatch(sku_id=None, similarity=0.0, rejected=True)
|
||||||
import numpy as np
|
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)
|
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
|
similarity = self._matrix @ embedding
|
||||||
# Several views per SKU: a SKU's score is its best view.
|
# Several views per SKU: a SKU's score is its best view.
|
||||||
best_by_sku: Dict[str, float] = {}
|
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,
|
return GalleryMatch(sku_id=top_sku, similarity=top_score, runner_up=runner_sku,
|
||||||
runner_up_similarity=runner_score)
|
runner_up_similarity=runner_score)
|
||||||
|
|
||||||
|
def head_loaded(self) -> bool:
|
||||||
|
return self._head is not None
|
||||||
|
|
||||||
def view_count(self) -> int:
|
def view_count(self) -> int:
|
||||||
return len(self._entries)
|
return len(self._entries)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -232,3 +232,54 @@ What is settled is the mechanism: **a head trained on labels the pipeline can pr
|
||||||
itself** - tag-named facings, plus operator confirmations - moves fine-grained
|
itself** - tag-named facings, plus operator confirmations - moves fine-grained
|
||||||
identity from useless to usable, and a rejection threshold turns it into something a
|
identity from useless to usable, and a rejection threshold turns it into something a
|
||||||
store can act on.
|
store can act on.
|
||||||
|
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
# Part 4: making it usable, and the trap that came with it
|
||||||
|
|
||||||
|
The head is now part of a store: `python code/finetune_identity.py --stock stock-lopia`
|
||||||
|
trains on that store's enrolled views and writes `<stock>/identity_head.pt`, which
|
||||||
|
`ProductGallery` loads automatically. The first attempt exposed a trap worth writing
|
||||||
|
down.
|
||||||
|
|
||||||
|
## A trained head destroys rejection
|
||||||
|
|
||||||
|
Measured on the same two sets (the home shelf it was enrolled from, and 20 shelves
|
||||||
|
whose products are absent):
|
||||||
|
|
||||||
|
| | home shelf accepted | foreign shelves accepted |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| raw features, threshold 0.55 | 78 % | **5 %** |
|
||||||
|
| trained head, threshold 0.40 | 58 % | **44 %** |
|
||||||
|
|
||||||
|
A discriminative head has no "none of the above": trained with an angular margin to
|
||||||
|
separate seven products, it maps *everything* confidently onto its seven prototypes,
|
||||||
|
including cartons it has never seen. Rejection, which was the gallery's whole point,
|
||||||
|
disappears.
|
||||||
|
|
||||||
|
## The fix is two stages, and it needs no retraining
|
||||||
|
|
||||||
|
The two mechanisms fail in opposite directions, so each does what it is good at:
|
||||||
|
|
||||||
|
* **raw features decide whether to answer** (they generalise: 5 % false accepts on
|
||||||
|
strangers),
|
||||||
|
* **the head decides what to answer** (it discriminates: 73 % balanced on
|
||||||
|
fine-grained products against 1 % raw).
|
||||||
|
|
||||||
|
| | home shelf accepted | foreign shelves accepted |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| **raw gate, then head ranking** | 47 % | **5 %** |
|
||||||
|
|
||||||
|
The home-shelf acceptance is lower than raw-only because the gate is strict, which is
|
||||||
|
the honest trade: fewer claims, and the claims are worth more.
|
||||||
|
|
||||||
|
## Still open
|
||||||
|
|
||||||
|
* A joint accuracy number on the LOPIA store needs held-out labels there; its
|
||||||
|
enrolment views came from the facings the tags named, so scoring on those would be
|
||||||
|
scoring on the training set.
|
||||||
|
* Training the head *with* out-of-distribution negatives (crops from other stores)
|
||||||
|
should let it keep rejecting on its own, which would remove the need for the gate.
|
||||||
|
* Everything above is still one photograph per store. The second visit remains the
|
||||||
|
measurement that settles it.
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue