diff --git a/app/models/megadetector.py b/app/models/megadetector.py index a66ef52..d797baa 100644 --- a/app/models/megadetector.py +++ b/app/models/megadetector.py @@ -9,6 +9,7 @@ from PytorchWildlife.models import detection as pw_detection from .base_model import BaseModel +from ..utils.checkpoint_utils import get_checkpoint_path logger = logging.getLogger(__name__) @@ -55,18 +56,27 @@ def load(self, model_path: str, device: str, **kwargs) -> None: device: Device to load the model on (e.g., 'cpu', 'cuda', 'mps') **kwargs: Additional parameters including: - model_id: The model variant (e.g., 'MDV6-yolov10-e') + - checkpoint_path: Optional explicit weight file (local path + or URL, resolved like every other model type). Absent -> + PytorchWildlife downloads its pinned weights itself. - conf: Default confidence threshold """ # Use model_id from kwargs as the version model_id = kwargs.get('model_id') if not model_id: raise ValueError("Model ID must be provided in the model configuration") - + + # Resolve an explicit weight through the shared resolver so + # MODEL_BASE-style configs can pin/serve this model like the others, + # instead of always downloading from PytorchWildlife's zenodo URL. + local_weights = get_checkpoint_path(kwargs.get('checkpoint_path')) + logger.info(f"Loading MegaDetector model {model_id} on device {device}") - + # Load the model using model_id as version try: self.model = pw_detection.MegaDetectorV6( + weights=local_weights, version=model_id, device=device, pretrained=True diff --git a/tests/test_megadetector_checkpoint.py b/tests/test_megadetector_checkpoint.py new file mode 100644 index 0000000..f205f8e --- /dev/null +++ b/tests/test_megadetector_checkpoint.py @@ -0,0 +1,81 @@ +"""Tests for explicit checkpoint_path support in the MegaDetector loader. + +Every other model type resolves its weights through +checkpoint_utils.get_checkpoint_path (local path or URL). MegaDetector +was the exception: it always let PytorchWildlife fetch weights from its +hardcoded zenodo URL at load time — a hidden internet download on every +cold start that MODEL_BASE cannot control or pin. + +PytorchWildlife needs GPU/network extras that aren't importable in the +test environment, so a stub module captures the constructor call. +""" +import sys +import types + +import pytest + + +@pytest.fixture +def megadetector_model(monkeypatch): + """Import app.models.megadetector against a stubbed PytorchWildlife.""" + captured = {} + + class FakeMegaDetectorV6: + def __init__(self, weights=None, device="cpu", pretrained=True, + version="yolov9c"): + captured.update(weights=weights, device=device, + pretrained=pretrained, version=version) + + detection = types.ModuleType("PytorchWildlife.models.detection") + detection.MegaDetectorV6 = FakeMegaDetectorV6 + models = types.ModuleType("PytorchWildlife.models") + models.detection = detection + pw = types.ModuleType("PytorchWildlife") + pw.models = models + + monkeypatch.setitem(sys.modules, "PytorchWildlife", pw) + monkeypatch.setitem(sys.modules, "PytorchWildlife.models", models) + monkeypatch.setitem(sys.modules, "PytorchWildlife.models.detection", detection) + monkeypatch.delitem(sys.modules, "app.models.megadetector", raising=False) + + from app.models.megadetector import MegaDetectorModel + yield MegaDetectorModel, captured + monkeypatch.delitem(sys.modules, "app.models.megadetector", raising=False) + + +def test_checkpoint_path_reaches_pytorchwildlife(megadetector_model, tmp_path): + MegaDetectorModel, captured = megadetector_model + weights = tmp_path / "MDV6-yolov10-e-1280.pt" + weights.write_bytes(b"w") + + model = MegaDetectorModel() + model.load(model_path="", device="cpu", model_id="MDV6-yolov10-e", + checkpoint_path=str(weights)) + + assert captured["weights"] == str(weights) + assert captured["version"] == "MDV6-yolov10-e" + + +def test_absent_checkpoint_path_keeps_auto_download(megadetector_model): + """Without checkpoint_path the current behavior is preserved: + PytorchWildlife receives weights=None and fetches its own.""" + MegaDetectorModel, captured = megadetector_model + + model = MegaDetectorModel() + model.load(model_path="", device="cpu", model_id="MDV6-yolov10-e", conf=0.1) + + assert captured["weights"] is None + assert captured["pretrained"] is True + + +def test_missing_local_checkpoint_fails_fast(megadetector_model): + """A configured-but-absent local weight is a config error, not a cue + to silently download something else.""" + MegaDetectorModel, captured = megadetector_model + + model = MegaDetectorModel() + with pytest.raises(Exception) as excinfo: + model.load(model_path="", device="cpu", model_id="MDV6-yolov10-e", + checkpoint_path="/nonexistent/weights.pt") + assert "not found" in str(excinfo.value).lower() + assert captured == {}, "PytorchWildlife must not be constructed on bad config"