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())},