fix: refuse unavailable gallery and pin head trainer

This commit is contained in:
kiyo9w 2026-09-28 14:41:57 +09:00
parent d5a7a56dff
commit ac315eb48e
4 changed files with 77 additions and 6 deletions

View file

@ -293,8 +293,11 @@ def create_app(
scans = scans if scans is not None else build_legacy_scans(engine)
# A live session can be recorded so a real scan becomes data to label later.
recorder = FrameRecorder(settings.record_frames, every=settings.record_every) \
if settings.record_frames else None
recorder = FrameRecorder(
settings.record_frames,
every=settings.record_every,
identity_backend=engine.settings.identity_backend,
) if settings.record_frames else None
samples = {item["id"]: item for item in get_available_shelf_images("add")}
# path -> (mtime_ns, size, width, height). Dropped when the file changes.
sample_dimensions: Dict[str, Tuple[int, int, int, int]] = {}
@ -1138,7 +1141,8 @@ def main() -> None:
engine = build_default_engine(engine_settings)
if (engine_settings.identity_backend or "").lower() == "gallery":
if engine.identity_gallery() is None:
raise RuntimeError(
engine.shutdown()
raise SystemExit(
"gallery identity backend requires local DINOv2 weights; refusing to start"
)
warm_ms = engine.warm_up()

View file

@ -132,12 +132,14 @@ def enrollment_report_text(report: dict) -> str:
def build_model(image_size: int = 518):
import timm
import torch
from gallery import load_dinov2
device = "mps" if torch.backends.mps.is_available() else "cpu"
model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
num_classes=0, img_size=image_size, dynamic_img_size=True)
# Same hash-pinned local file as inference. pretrained=True would fetch a
# different copy even with YOLO_OFFLINE=1.
model = load_dinov2(image_size)
return model.eval().to(device), device

View file

@ -274,6 +274,33 @@ class CatalogTests(unittest.TestCase):
self.assertEqual(product["number_out"], 3)
self.assertIsNone(engine.update_stock("nope", -1))
class GalleryLoadTests(unittest.TestCase):
def test_a_failed_gallery_load_returns_none_and_is_not_retried(self):
engine, _, _, _ = make_engine(identity_backend="gallery", stock_dir="does-not-exist")
class BrokenGallery:
def __init__(self, *args, **kwargs):
pass
def load(self):
raise FileNotFoundError("DINOv2 weights not found")
import gallery as gallery_module
original = gallery_module.ProductGallery
gallery_module.ProductGallery = BrokenGallery
try:
self.assertIsNone(engine.identity_gallery())
self.assertIsNone(engine.identity_gallery())
finally:
gallery_module.ProductGallery = original
engine.shutdown()
self.assertIs(engine._gallery, False)
class CatalogImageTests(unittest.TestCase):
def test_catalog_image_path_is_resolved_through_inventory(self):
engine, _, _, _ = make_engine()
self.assertIsNone(engine.catalog_image_path("../../etc/passwd"))

View file

@ -225,5 +225,43 @@ class HeadPinTests(unittest.TestCase):
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()