fix: refuse unavailable gallery and pin head trainer
This commit is contained in:
parent
d5a7a56dff
commit
ac315eb48e
4 changed files with 77 additions and 6 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue