148 lines
5.7 KiB
Python
148 lines
5.7 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_SHA256,
|
|
GalleryEntry,
|
|
ProductGallery,
|
|
dinov2_weights_path,
|
|
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
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|