tngo/product-recognition-demo
1
0
Fork 0
forked from vnguyen/product-recognition-demo
Code Pull requests Activity

fix: constrain trained head labels to enrolled catalog

Browse source
This commit is contained in:
kiyo9w 2026-09-28 14:07:02 +09:00
parent b937be8e30
commit 2798322417
3 changed files with 184 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

9
recognization/code/finetune_identity.py
View file

@ -228,7 +228,7 @@ def train_from_stock(args) -> None:
import torch
from gallery import ProductGallery
from gallery import DINOV2_MODEL, ProductGallery, catalog_digest, head_class_set
stock = args.stock
report = enrollment_report(stock)
@ -267,10 +267,15 @@ def train_from_stock(args) -> None:
projection, prototypes = train_head(features, augmented_labels, class_ids,
args.epochs, device, margin=0.25)
digest = catalog_digest(stock, class_ids)
if head_class_set(class_ids) is None or digest is None:
print(f"{stock}: refusing to write a head; class ids are empty or duplicated")
return
payload = {
"class_ids": class_ids,
"catalog_digest": digest,
"image_size": 518,
"base_model": "vit_small_patch14_dinov2.lvd142m",
"base_model": DINOV2_MODEL,
"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())},

117
recognization/code/gallery.py
View file

@ -25,6 +25,7 @@ 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
@ -131,6 +132,39 @@ class GalleryMatch:
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."""
@ -162,45 +196,86 @@ class ProductGallery:
# 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, if it has one.
"""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 payload that is not that head leaves the raw gallery in place.
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)
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],
}
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

81
recognization/code/test_gallery.py
View file

@ -22,10 +22,13 @@ if current_dir not in sys.path:
sys.path.insert(0, current_dir)
from gallery import ( # noqa: E402
DINOV2_MODEL,
DINOV2_SHA256,
GalleryEntry,
ProductGallery,
catalog_digest,
dinov2_weights_path,
head_class_set,
load_dinov2,
)
@ -144,5 +147,83 @@ class OfflineWeightsTests(unittest.TestCase):
if previous is not None:
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
def _head_payload(class_ids, torch):
projection = torch.nn.Linear(384, 128)
prototypes = torch.nn.functional.normalize(torch.randn(len(class_ids), 128), dim=1)
return {
"class_ids": list(class_ids),
"image_size": 518,
"base_model": DINOV2_MODEL,
"projection": {key: value.detach().cpu() for key, value in projection.state_dict().items()},
"prototypes": prototypes.detach().cpu(),
}
class HeadPinTests(unittest.TestCase):
def test_a_stale_head_cannot_label_this_catalog(self):
import tempfile
import torch
with tempfile.TemporaryDirectory() as directory:
with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
handle.write("SKU_ID,product_name\nsku_a,A\nsku_b,B\nsku_c,C\n")
gallery = ProductGallery(directory, ["sku_a", "sku_b", "sku_c"])
gallery._available = True
path = gallery._head_path
subset = ["sku_c", "sku_a"]
digest = catalog_digest(directory, subset)
foreign = _head_payload(["other_store", "sku_a"], torch)
torch.save(foreign, path)
self.assertFalse(gallery.load_head())
self.assertIsNone(gallery._head)
self.assertIn("catalog_mismatch", gallery.head_reason)
self.assertTrue(gallery.available())
legacy = _head_payload(subset, torch)
torch.save(legacy, path)
self.assertTrue(gallery.load_head())
self.assertEqual(gallery.head_reason, "legacy_unpinned")
self.assertEqual(gallery._head["class_ids"], subset)
self.assertNotIn("other_store", gallery._head["class_ids"])
stale = _head_payload(subset, torch)
stale["catalog_digest"] = "0" * 64
torch.save(stale, path)
self.assertFalse(gallery.load_head())
self.assertIsNone(gallery._head)
self.assertIn("catalog_mismatch", gallery.head_reason)
self.assertTrue(gallery.available())
wrong_backbone = _head_payload(subset, torch)
wrong_backbone["catalog_digest"] = digest
wrong_backbone["base_model"] = "vit_base_patch14_dinov2.lvd142m"
torch.save(wrong_backbone, path)
self.assertFalse(gallery.load_head())
self.assertIn("backbone_mismatch", gallery.head_reason)
wrong_size = _head_payload(subset, torch)
wrong_size["catalog_digest"] = digest
wrong_size["image_size"] = 224
torch.save(wrong_size, path)
self.assertFalse(gallery.load_head())
self.assertIn("image_size_mismatch", gallery.head_reason)
torch.save(["not", "a", "head"], path)
self.assertFalse(gallery.load_head())
self.assertTrue(gallery.head_reason.startswith("malformed"))
self.assertTrue(gallery.available())
self.assertIsNone(gallery._head)
current = _head_payload(subset, torch)
current["catalog_digest"] = digest
torch.save(current, path)
self.assertTrue(gallery.load_head())
self.assertEqual(gallery.head_reason, "ok")
self.assertEqual(gallery._head["class_ids"], subset)
self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))
if __name__ == "__main__":
unittest.main()