02e69acac7
A cache key proves the store, not the rows inside the file. The old check asked only that every cached row appear somewhere in the enrolled list, so a row moved across SKUs still passed: the matrix that travelled with it would then be filed under the wrong product, and a valid key would rename a SKU. Rows now walk the enrolled list forward and never backwards, so a moved, repeated or foreign row has nowhere left to match. Views that could not be decoded stay absent without disturbing the order of the ones after them, which is what keeps a shorter cache a hit rather than a re-embed.
575 lines
25 KiB
Python
575 lines
25 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,
|
|
)
|
|
|
|
import gallery as gallery_module # noqa: E402
|
|
|
|
|
|
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 Backbone:
|
|
"""A deterministic stand-in for DINOv2: the prepared crop's channel means,
|
|
tiled to the width the real backbone returns."""
|
|
|
|
def __init__(self):
|
|
self.calls = 0
|
|
|
|
def __call__(self, tensor):
|
|
self.calls += 1
|
|
return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3)
|
|
|
|
|
|
def a_query():
|
|
"""A unit vector in the backbone's space, for match() without a model."""
|
|
query = np.zeros(gallery_module.DINOV2_EMBED_DIM, dtype=np.float32)
|
|
query[0] = 1.0
|
|
query[1] = 0.05
|
|
return query / np.linalg.norm(query)
|
|
|
|
|
|
class UnreachableBackbone:
|
|
def __call__(self, tensor):
|
|
raise AssertionError("the gallery loaded a backbone it should have cached")
|
|
|
|
|
|
def write_view(path, seed):
|
|
import cv2
|
|
|
|
image = np.random.RandomState(seed).randint(0, 256, (24, 32, 3), dtype=np.uint8)
|
|
cv2.imwrite(path, image)
|
|
|
|
|
|
def a_store(directory, names=("sku_a", "sku_a_1", "sku_b")):
|
|
"""A store dir: one JPEG per view plus the stock.csv the catalog is read from."""
|
|
with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
|
|
handle.write("SKU_ID,product_name\n")
|
|
for index, name in enumerate(names):
|
|
write_view(os.path.join(directory, f"{name}.jpeg"), index)
|
|
|
|
|
|
def a_stored_gallery(directory, names, backbone, image_size=32):
|
|
gallery = ProductGallery(directory, list(names), device="cpu", image_size=image_size)
|
|
gallery._load_model = lambda: backbone
|
|
return gallery
|
|
|
|
|
|
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_the_per_sku_max_does_not_depend_on_entry_order(self):
|
|
# Three SKUs, several views each, enrolled out of order: a SKU is scored
|
|
# by its best view, and the answer cannot move because a row moved.
|
|
gallery = a_gallery([
|
|
("sku_b", unit([0, 1, 0])),
|
|
("sku_a", unit([1, 0, 0])),
|
|
("sku_c", unit([0, 0, 1])),
|
|
("sku_b", unit([0.3, 0.95, 0])),
|
|
("sku_a", unit([0.98, 0.2, 0])),
|
|
])
|
|
gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32)
|
|
straight = gallery.match(object())
|
|
self.assertEqual(straight.sku_id, "sku_a")
|
|
self.assertEqual(straight.runner_up, "sku_b")
|
|
|
|
for order in ([4, 0, 3, 1, 2], [2, 3, 1, 4, 0]):
|
|
gallery._entries = [gallery._entries[index] for index in order]
|
|
gallery._matrix = gallery._matrix[order]
|
|
shuffled = gallery.match(object())
|
|
self.assertEqual(shuffled.sku_id, straight.sku_id)
|
|
self.assertEqual(shuffled.runner_up, straight.runner_up)
|
|
self.assertAlmostEqual(shuffled.similarity, straight.similarity, places=6)
|
|
self.assertAlmostEqual(shuffled.runner_up_similarity, straight.runner_up_similarity, places=6)
|
|
|
|
def test_a_tie_keeps_the_first_enrolled_sku(self):
|
|
# Exactly equidistant: first appearance decides, as it did before the
|
|
# per-SKU maximum became a segmented reduction.
|
|
gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
|
|
gallery.embed = lambda crop: np.array([1.0, 1.0, 0.0], dtype=np.float32)
|
|
match = gallery.match(object())
|
|
self.assertFalse(match.rejected)
|
|
self.assertEqual(match.sku_id, "sku_a")
|
|
self.assertEqual(match.runner_up, "sku_b")
|
|
self.assertAlmostEqual(match.similarity, match.runner_up_similarity, places=6)
|
|
self.assertAlmostEqual(match.margin, 0.0, places=6)
|
|
|
|
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 FeatureCacheTests(unittest.TestCase):
|
|
"""The cached matrix must be this store's, or it must not be used."""
|
|
|
|
def setUp(self):
|
|
import hashlib
|
|
import tempfile
|
|
|
|
self.directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.directory.cleanup)
|
|
a_store(self.directory.name)
|
|
# A cache is only served when the pinned backbone is on disk, so point the
|
|
# pin at a stand-in file the module's expected digest is told about. The
|
|
# check itself - hash this file, compare - is what is under test.
|
|
self.weights = os.path.join(self.directory.name, "dinov2.safetensors")
|
|
self.write(self.weights, b"not really a checkpoint")
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.weights
|
|
self.addCleanup(os.environ.pop, "DORAYAKI_DINOV2_WEIGHTS", None)
|
|
self.addCleanup(setattr, gallery_module, "DINOV2_SHA256", gallery_module.DINOV2_SHA256)
|
|
gallery_module.DINOV2_SHA256 = hashlib.sha256(b"not really a checkpoint").hexdigest()
|
|
|
|
def write(self, path, content):
|
|
with open(path, "wb") as handle:
|
|
handle.write(content)
|
|
return path
|
|
|
|
def real_gallery(self):
|
|
"""A gallery with the real _load_model: the server's own readiness path."""
|
|
return ProductGallery(self.directory.name, ["sku_a", "sku_b"], device="cpu", image_size=32)
|
|
|
|
def loaded(self, names=("sku_a", "sku_b"), image_size=32):
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, names, backbone, image_size=image_size)
|
|
self.assertTrue(gallery.load())
|
|
return gallery, backbone
|
|
|
|
def cache_payload(self, gallery):
|
|
"""The rows and matrix a loaded gallery left on disk."""
|
|
with np.load(gallery._cache_path(), allow_pickle=False) as payload:
|
|
return (
|
|
[str(value) for value in payload["skus"].tolist()],
|
|
[str(value) for value in payload["views"].tolist()],
|
|
payload["matrix"],
|
|
)
|
|
|
|
def rewrite_cache(self, gallery, skus, views, matrix):
|
|
"""The same valid key with different contents: only the rows can catch it."""
|
|
path = gallery._cache_path()
|
|
with np.load(path, allow_pickle=False) as payload:
|
|
key = str(payload["key"].item())
|
|
np.savez(path, key=key, skus=np.array(skus), views=np.array(views), matrix=matrix)
|
|
|
|
def re_embeds(self, first, why):
|
|
"""The one right answer to a rewritten cache: build it from the pixels."""
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(gallery.load(), why)
|
|
self.assertEqual(backbone.calls, 1, f"{why} must not be served from cache")
|
|
np.testing.assert_array_equal(gallery._matrix, first._matrix)
|
|
self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_a", "sku_b"])
|
|
|
|
def test_a_second_start_answers_the_same_from_the_cache(self):
|
|
first, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertEqual(first.view_count(), 3)
|
|
self.assertTrue(os.path.isfile(first._cache_path()))
|
|
|
|
cached = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone())
|
|
self.assertTrue(cached.load())
|
|
self.assertEqual(cached.view_count(), first.view_count())
|
|
np.testing.assert_array_equal(cached._matrix, first._matrix)
|
|
self.assertEqual([entry.view for entry in cached._entries],
|
|
[entry.view for entry in first._entries])
|
|
query = a_query()
|
|
first.embed = cached.embed = lambda crop: query
|
|
self.assertEqual(vars(cached.match(object())), vars(first.match(object())))
|
|
|
|
def test_changed_view_bytes_re_embed(self):
|
|
first, _ = self.loaded()
|
|
before = first._matrix.copy()
|
|
write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)
|
|
|
|
second, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1, "a replaced view must not be served from cache")
|
|
self.assertFalse(np.array_equal(second._matrix, before))
|
|
|
|
def test_a_reordered_catalog_does_not_borrow_another_mapping(self):
|
|
self.loaded()
|
|
swapped, backbone = self.loaded(names=("sku_b", "sku_a"))
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertEqual([entry.sku_id for entry in swapped._entries], ["sku_b", "sku_a", "sku_a"])
|
|
|
|
def test_another_image_size_is_never_served(self):
|
|
small, _ = self.loaded(image_size=32)
|
|
large, backbone = self.loaded(image_size=64)
|
|
self.assertEqual(backbone.calls, 1)
|
|
self.assertNotEqual(small._cache_path(), large._cache_path())
|
|
|
|
def test_a_corrupt_cache_re_embeds_the_same_matrix(self):
|
|
first, _ = self.loaded()
|
|
with open(first._cache_path(), "wb") as handle:
|
|
handle.write(b"not an npz at all")
|
|
|
|
second, backbone = self.loaded()
|
|
self.assertEqual(backbone.calls, 1)
|
|
np.testing.assert_array_equal(second._matrix, first._matrix)
|
|
|
|
def test_swapped_cache_rows_are_not_served(self):
|
|
# A current key says nothing about the rows inside the file. The two
|
|
# sku_a views are adjacent, so a swap has to cross SKUs to mean
|
|
# anything - and the matrix that moves with them is the dangerous part:
|
|
# a vector that would then be filed under the wrong product.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.assertEqual(list(zip(skus, views)),
|
|
[("sku_a", "sku_a.jpeg"), ("sku_a", "sku_a_1.jpeg"), ("sku_b", "sku_b.jpeg")])
|
|
self.rewrite_cache(first, skus[::-1], views[::-1], matrix[::-1])
|
|
self.re_embeds(first, "rows swapped with their vectors")
|
|
|
|
def test_mislabelled_cache_rows_are_not_served(self):
|
|
# The views and the matrix stay where they are, so every row names a
|
|
# product it was not embedded for.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, [skus[2], skus[1], skus[0]], views, matrix)
|
|
self.re_embeds(first, "labels swapped around the views")
|
|
|
|
def test_a_duplicated_cache_row_is_not_served(self):
|
|
# The first view written twice: one image would answer for two products.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[:1] * 2 + skus[2:], views[:1] * 2 + views[2:], matrix)
|
|
self.re_embeds(first, "a row repeated")
|
|
|
|
def test_a_foreign_sku_row_is_not_served(self):
|
|
# A row from another store's catalog: a SKU this gallery never enrolled.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[1:] + ["other_store"], views[1:] + ["other_store.jpeg"], matrix)
|
|
self.re_embeds(first, "another store's row")
|
|
|
|
def test_a_cache_may_omit_a_view_it_could_not_decode(self):
|
|
# The writer drops a view it could not read and keeps everyone else's
|
|
# order, so such a cache is still this store's - and still worth serving
|
|
# rather than re-embedding the whole catalog on every start.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
self.rewrite_cache(first, skus[:1] + skus[2:], views[:1] + views[2:], matrix[[0, 2]])
|
|
|
|
backbone = Backbone()
|
|
gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 0, "a readable shorter cache is still a hit")
|
|
self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_b"])
|
|
np.testing.assert_array_equal(gallery._matrix, first._matrix[[0, 2]])
|
|
|
|
def test_a_nonnumeric_cache_matrix_is_not_served(self):
|
|
# Only the float matrix _cache_write produces is a cache. A complex
|
|
# matrix is finite and casts without complaint, so it answers with a
|
|
# different matrix than the one on disk; integer, boolean and text
|
|
# matrices are not embeddings at all.
|
|
first, _ = self.loaded()
|
|
skus, views, matrix = self.cache_payload(first)
|
|
others = {
|
|
"complex": matrix.astype(np.complex128),
|
|
"timedelta": np.zeros_like(matrix, dtype="timedelta64[s]"),
|
|
"integer": (matrix * 8).astype(np.int32),
|
|
"boolean": np.ones_like(matrix, dtype=bool),
|
|
"text": np.zeros(matrix.shape, dtype="<U4"),
|
|
}
|
|
for name, other in others.items():
|
|
with self.subTest(matrix=name):
|
|
self.rewrite_cache(first, skus, views, other)
|
|
self.re_embeds(first, f"a {name} matrix")
|
|
|
|
def test_a_loaded_gallery_reloads_when_its_source_moves(self):
|
|
gallery, backbone = self.loaded()
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 1, "an unchanged store must not re-embed")
|
|
|
|
write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 2)
|
|
|
|
with open(os.path.join(self.directory.name, "stock.csv"), "a", encoding="utf-8") as handle:
|
|
handle.write("sku_c,new product\n")
|
|
self.assertTrue(gallery.load())
|
|
self.assertEqual(backbone.calls, 3, "stock.csv is part of what an embedding is pinned to")
|
|
self.assertEqual(gallery.view_count(), 3)
|
|
|
|
def test_a_valid_cache_needs_the_pinned_backbone_on_disk(self):
|
|
# The regression: a cache hit never opens load_dinov2, so a deleted or
|
|
# altered checkpoint used to read as a loaded gallery and then raise on
|
|
# the first match - after the server had reported itself ready.
|
|
self.loaded()
|
|
gallery = self.real_gallery()
|
|
self.assertTrue(gallery.load(), "the cache serves while the pin holds")
|
|
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(self.directory.name, "gone.safetensors")
|
|
missing = self.real_gallery()
|
|
self.assertFalse(missing.load())
|
|
self.assertFalse(missing.available())
|
|
|
|
os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.write(
|
|
os.path.join(self.directory.name, "altered.safetensors"), b"a different checkpoint"
|
|
)
|
|
altered = self.real_gallery()
|
|
self.assertFalse(altered.load())
|
|
self.assertFalse(altered.available())
|
|
|
|
def test_an_altered_backbone_re_embeds_rather_than_serves(self):
|
|
# With a backbone that can still load, an altered file must re-embed.
|
|
self.loaded()
|
|
self.assertTrue(
|
|
a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone()).load(),
|
|
"the cache still serves on the pinned file",
|
|
)
|
|
self.write(self.weights, b"a different checkpoint")
|
|
backbone = Backbone()
|
|
rebuilt = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
|
|
self.assertTrue(rebuilt.load())
|
|
self.assertEqual(backbone.calls, 1, "an unpinned cache must not be served")
|
|
|
|
|
|
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
|
|
|
|
|
|
# build_model resolves the loader from gallery at call time, so that is
|
|
# the attribute to replace; finetune_identity has no module-level copy.
|
|
original = gallery_module.load_dinov2
|
|
gallery_module.load_dinov2 = fake_load
|
|
try:
|
|
model, device = finetune_identity.build_model(224)
|
|
finally:
|
|
gallery_module.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()
|