From 0e85c399bbc4c6c661eb3dd7a07a2b8655ddbf3f Mon Sep 17 00:00:00 2001 From: kiyo9w Date: Mon, 28 Sep 2026 13:41:27 +0900 Subject: [PATCH] fix: load hash-pinned DINOv2 weights without network --- recognization/.gitignore | 1 + recognization/code/gallery.py | 80 +++++++++++++++++++++++++++--- recognization/code/test_gallery.py | 38 +++++++++++++- 3 files changed, 111 insertions(+), 8 deletions(-) diff --git a/recognization/.gitignore b/recognization/.gitignore index a0bb35a..013e35c 100644 --- a/recognization/.gitignore +++ b/recognization/.gitignore @@ -19,6 +19,7 @@ ENV/ # Temporary weights runs/detect/train/weights/last.pt +weights/dinov2_vits14_lvd142m.safetensors *.log stock-*/identity_head.pt .demo-token diff --git a/recognization/code/gallery.py b/recognization/code/gallery.py index 0805d71..5d5f5dd 100644 --- a/recognization/code/gallery.py +++ b/recognization/code/gallery.py @@ -38,6 +38,78 @@ DEFAULT_THRESHOLD = 0.55 # The same line for a fine-tuned head, whose similarity scale is its own. DEFAULT_HEAD_THRESHOLD = 0.40 +# DINOv2 ViT-S/14, LVD-142M. timm's default source is the Hugging Face copy of +# these weights (Apache-2.0); pretrained=True fetches it even when YOLO_OFFLINE=1. +# The pin is the sha256 of that model.safetensors (timm hub snapshot 4610ca14). +# Not in Git (~84 MB). Place it at weights/dinov2_vits14_lvd142m.safetensors, or +# point DORAYAKI_DINOV2_WEIGHTS at a copy. See README "DINOv2 weights". +DINOV2_MODEL = "vit_small_patch14_dinov2.lvd142m" +DINOV2_SHA256 = "04d27f3400d059fc0cfd7d17dd1909a75bf3ea8fb3eeb48b97cb99e57ee20081" +DINOV2_FILENAME = "dinov2_vits14_lvd142m.safetensors" + + +def dinov2_weights_path() -> str: + override = os.environ.get("DORAYAKI_DINOV2_WEIGHTS") + if override: + return override + here = os.path.dirname(os.path.abspath(__file__)) + return os.path.join(os.path.dirname(here), "weights", DINOV2_FILENAME) + + +def _sha256(path: str) -> str: + import hashlib + + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def load_dinov2(image_size: int) -> Any: + """Build the gallery backbone from the pinned local file. Never downloads. + + A missing or altered file raises FileNotFoundError naming the path and the + expected sha256. The head in finetune_identity.py is a Linear(384, 128) on + this model's features; a different backbone will not load that head. + """ + import timm + from safetensors.torch import load_file + from timm.models.vision_transformer import checkpoint_filter_fn + + path = dinov2_weights_path() + if not os.path.isfile(path): + raise FileNotFoundError( + f"DINOv2 weights not found at {path}. " + f"Copy model.safetensors (sha256 {DINOV2_SHA256}) from the local " + f"Hugging Face cache of timm/{DINOV2_MODEL}, or set DORAYAKI_DINOV2_WEIGHTS. " + f"The gallery does not download weights." + ) + digest = _sha256(path) + if digest != DINOV2_SHA256: + raise FileNotFoundError( + f"DINOv2 weights at {path} are sha256 {digest}, expected {DINOV2_SHA256}." + ) + model = timm.create_model( + DINOV2_MODEL, + pretrained=False, + num_classes=0, + img_size=int(image_size), + dynamic_img_size=True, + ) + # The file is the 518px checkpoint. timm's own filter resizes pos_embed when + # the caller asks for another size; pretrained=True would do this and also + # download, so it stays off. + state = checkpoint_filter_fn(load_file(path), model) + missing, unexpected = model.load_state_dict(state, strict=False) + unexpected = [key for key in unexpected if not key.startswith("head.")] + if missing or unexpected: + raise RuntimeError( + f"DINOv2 weights at {path} do not match {DINOV2_MODEL} " + f"(missing={missing}, unexpected={unexpected})." + ) + return model + @dataclass class GalleryEntry: @@ -138,13 +210,7 @@ class ProductGallery: device = self.device if device is None: device = "mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu" - model = timm.create_model( - "vit_small_patch14_dinov2.lvd142m", - pretrained=True, - num_classes=0, - img_size=self.image_size, - dynamic_img_size=True, - ) + model = load_dinov2(self.image_size) model = model.eval().to(device) self.device = device self._model = model diff --git a/recognization/code/test_gallery.py b/recognization/code/test_gallery.py index afc6585..815365e 100644 --- a/recognization/code/test_gallery.py +++ b/recognization/code/test_gallery.py @@ -21,7 +21,13 @@ if os.path.exists(os.path.join(project_root, "stock")): if current_dir not in sys.path: sys.path.insert(0, current_dir) -from gallery import GalleryEntry, ProductGallery # noqa: E402 +from gallery import ( # noqa: E402 + DINOV2_SHA256, + GalleryEntry, + ProductGallery, + dinov2_weights_path, + load_dinov2, +) def a_gallery(views, threshold=0.55): @@ -108,5 +114,35 @@ class EnrollmentTests(unittest.TestCase): self.assertEqual(len(gallery._views_for("sku_ab")), 1) + +class OfflineWeightsTests(unittest.TestCase): + def test_a_missing_weights_file_names_the_path_and_does_not_fetch(self): + previous = os.environ.get("DORAYAKI_DINOV2_WEIGHTS") + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join( + "does-not-exist", "dinov2_vits14_lvd142m.safetensors" + ) + try: + with self.assertRaises(FileNotFoundError) as raised: + load_dinov2(224) + finally: + if previous is None: + os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None) + else: + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous + message = str(raised.exception) + self.assertIn("does-not-exist", message) + self.assertIn(DINOV2_SHA256, message) + self.assertIn("does not download", message) + + def test_the_default_path_stays_inside_the_repo(self): + previous = os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None) + try: + self.assertTrue(dinov2_weights_path().endswith( + os.path.join("weights", "dinov2_vits14_lvd142m.safetensors") + )) + finally: + if previous is not None: + os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous + if __name__ == "__main__": unittest.main()