tngo/product-recognition-demo
1
0
Fork 0
forked from vnguyen/product-recognition-demo
Code Pull requests Activity

feat: add bounded crop evidence and keyed gallery cache

Browse source
This commit is contained in:
kiyo9w 2026-09-28 15:49:34 +09:00
parent 69de9ea418
commit 0d0e9f82b9
8 changed files with 1065 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

250
recognization/code/test_gallery.py
View file

@ -32,6 +32,8 @@ from gallery import ( # noqa: E402
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."""
@ -55,6 +57,52 @@ def unit(vector):
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]))])
@ -98,6 +146,42 @@ class MatchingTests(unittest.TestCase):
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())
@ -117,6 +201,164 @@ class EnrollmentTests(unittest.TestCase):
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 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_a_tampered_cache_row_is_not_served(self):
# A current key says nothing about the rows inside the file. Swap them
# and the gallery must re-embed, never name a SKU by a moved row.
first, _ = self.loaded()
path = first._cache_path()
with np.load(path, allow_pickle=False) as payload:
key = str(payload["key"].item())
skus = [str(value) for value in payload["skus"].tolist()]
views = [str(value) for value in payload["views"].tolist()]
matrix = payload["matrix"]
# The two sku_a views are adjacent, so a label swap has to cross SKUs to
# mean anything - and the matrix that stays put is the dangerous one:
# a vector that would then be filed under the wrong product.
self.assertEqual(list(zip(skus, views)),
[("sku_a", "sku_a.jpeg"), ("sku_a", "sku_a_1.jpeg"), ("sku_b", "sku_b.jpeg")])
tampered = {
"reordered rows": (skus[::-1], views[::-1], matrix[::-1]),
"swapped labels": ([skus[2], skus[1], skus[0]], views, matrix),
"another store's row": (skus[1:] + ["other_store"], views + ["other_store.jpeg"], matrix),
}
for label, (bad_skus, bad_views, bad_matrix) in tampered.items():
np.savez(path, key=key, skus=np.array(bad_skus), views=np.array(bad_views), matrix=bad_matrix)
backbone = Backbone()
gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
self.assertTrue(gallery.load(), label)
self.assertEqual(backbone.calls, 1, f"{label} 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_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):
@ -246,12 +488,14 @@ class TrainerPinTests(unittest.TestCase):
return sentinel
original = finetune_identity.load_dinov2
finetune_identity.load_dinov2 = fake_load
# 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:
finetune_identity.load_dinov2 = original
gallery_module.load_dinov2 = original
self.assertIs(model, sentinel)
self.assertIn(device, ("cpu", "mps"))
self.assertEqual(seen, [224])