""" Tests for the product gallery's matching logic. The embedding model is injected, so these run without torch or a GPU: what is being tested is the retrieval contract - best view wins, several views per SKU count as one answer, and a stranger is rejected rather than named. """ from __future__ import annotations import os import sys import unittest import numpy as np current_dir = os.path.dirname(os.path.abspath(__file__)) project_root = os.path.dirname(current_dir) if os.path.exists(os.path.join(project_root, "stock")): os.chdir(project_root) 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, ) def a_gallery(views, threshold=0.55): """A gallery with hand-made embeddings: no model, no files.""" skus = sorted({sku for sku, _ in views}) gallery = ProductGallery("does-not-exist", skus, threshold=threshold) embeddings = [] entries = [] for index, (sku, vector) in enumerate(views): vector = np.asarray(vector, dtype=np.float32) vector /= np.linalg.norm(vector) embeddings.append(vector) entries.append(GalleryEntry(sku_id=sku, view=f"{sku}_{index}", embedding=vector)) gallery._entries = entries gallery._matrix = np.stack(embeddings) gallery._available = True return gallery def unit(vector): vector = np.asarray(vector, dtype=np.float32) return (vector / np.linalg.norm(vector)).tolist() class MatchingTests(unittest.TestCase): def test_the_nearest_view_names_the_product(self): gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))]) gallery.embed = lambda crop: np.array([0.9, 0.1, 0.0], dtype=np.float32) match = gallery.match(object()) self.assertFalse(match.rejected) self.assertEqual(match.sku_id, "sku_a") self.assertAlmostEqual(match.similarity, 0.9, places=3) def test_several_views_of_one_sku_are_one_answer(self): # The second view of sku_b is closer than any view of sku_a, but sku_a must # still win on its best view: a SKU is not outvoted by having more photos. gallery = a_gallery([ ("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0])), ("sku_b", unit([0.2, 0.98, 0])), ]) gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32) match = gallery.match(object()) self.assertEqual(match.sku_id, "sku_a") self.assertEqual(match.runner_up, "sku_b") def test_a_stranger_is_rejected_not_named(self): gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))]) gallery.embed = lambda crop: np.array([0.0, 0.0, 1.0], dtype=np.float32) match = gallery.match(object()) self.assertTrue(match.rejected) self.assertIsNone(match.sku_id) def test_the_threshold_moves_the_rejection_line(self): gallery = a_gallery([("sku_a", unit([1, 0, 0]))], threshold=0.99) gallery.embed = lambda crop: np.array([0.9, 0.4, 0.0], dtype=np.float32) self.assertTrue(gallery.match(object()).rejected) self.assertFalse(gallery.match(object(), threshold=0.5).rejected) def test_margin_separates_a_confident_match_from_a_close_one(self): gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0.9, 0.44, 0]))]) gallery.embed = lambda crop: np.array([1.0, 0.1, 0.0], dtype=np.float32) confident = gallery.match(object()) gallery.embed = lambda crop: np.array([0.95, 0.3, 0.0], dtype=np.float32) close = gallery.match(object()) self.assertGreater(confident.margin, close.margin) def test_an_unloaded_gallery_rejects_everything(self): gallery = ProductGallery("does-not-exist", ["sku_a"]) match = gallery.match(object()) self.assertTrue(match.rejected) class EnrollmentTests(unittest.TestCase): def test_views_are_found_by_sku_prefix(self): import tempfile with tempfile.TemporaryDirectory() as directory: for name in ("sku_a.jpeg", "sku_a_1.jpeg", "sku_a-2.jpeg", "sku_ab.jpeg", "notes.txt"): open(os.path.join(directory, name), "wb").close() gallery = ProductGallery(directory, ["sku_a", "sku_ab"]) self.assertEqual(len(gallery._views_for("sku_a")), 3, "one view, two extra views, and not sku_ab") 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 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", ""])) class TrainerPinTests(unittest.TestCase): def test_head_training_loads_the_pinned_file_and_keeps_preprocess(self): import finetune_identity import gallery as gallery_module seen = [] class FakeBackbone: def eval(self): return self def to(self, device): return self sentinel = FakeBackbone() def fake_load(image_size): seen.append(image_size) return sentinel original = finetune_identity.load_dinov2 finetune_identity.load_dinov2 = fake_load try: model, device = finetune_identity.build_model(224) finally: finetune_identity.load_dinov2 = original self.assertIs(model, sentinel) self.assertIn(device, ("cpu", "mps")) self.assertEqual(seen, [224]) image = np.zeros((10, 20, 3), dtype=np.uint8) self.assertEqual(finetune_identity.prepare(image, 32).shape, gallery_module._prepare(image, 32).shape) self.assertTrue(np.allclose( finetune_identity.prepare(image, 32), gallery_module._prepare(image, 32), )) if __name__ == "__main__": unittest.main()