229 lines
9.1 KiB
Python
229 lines
9.1 KiB
Python
"""
|
|
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", ""]))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|