feat(models): selector dispatch for volumetric vs 2D MRI models

src/models/mri_selector.py

@@ -0,0 +1,48 @@
+"""Env-var-driven dispatch between volumetric ONNX and 2D resnet18 MRI models."""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+from src.core.logger import get_logger
+from src.models import mri_dl_2d, mri_model
+
+logger = get_logger(__name__)
+
+VALID_KINDS = ("volumetric_onnx", "resnet18_2d")
+_DEFAULT_KIND = "volumetric_onnx"
+
+
+def current_kind() -> str:
+    kind = os.environ.get("MRI_MODEL_KIND", _DEFAULT_KIND)
+    if kind not in VALID_KINDS:
+        raise ValueError(f"unknown MRI_MODEL_KIND={kind!r}; expected one of {VALID_KINDS}")
+    return kind
+
+
+def label_names_for_kind(kind: str) -> tuple[str, ...]:
+    if kind == "resnet18_2d":
+        return tuple(mri_dl_2d.IDX_TO_CLASS[i] for i in range(len(mri_dl_2d.CLASS_TO_IDX)))
+    return mri_model.DEFAULT_LABEL_NAMES
+
+
+def predict(
+    input_path: Path,
+    checkpoint_path: Path,
+    target_shape: tuple[int, int, int] | None = None,
+    label_names: tuple[str, ...] | None = None,
+) -> dict[str, Any]:
+    """Run the active MRI model on one input. Returns the unified prediction dict."""
+    kind = current_kind()
+    logger.info("dispatching MRI prediction kind=%s input=%s", kind, input_path)
+    if kind == "resnet18_2d":
+        model = mri_dl_2d.load(checkpoint_path)
+        return mri_dl_2d.predict_image(model, input_path)
+    model = mri_model.load(checkpoint_path)
+    return mri_model.predict_nifti(
+        model,
+        input_path,
+        target_shape=target_shape or mri_model.DEFAULT_TARGET_SHAPE,
+        label_names=label_names,
+    )

tests/models/test_mri_selector.py

@@ -0,0 +1,52 @@
+"""Tests for src.models.mri_selector — env-var-driven 2D / 3D dispatch."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from src.models import mri_selector
+from tests.fixtures.build_dummy_mri_onnx import build as build_dummy_3d
+from tests.fixtures.build_dummy_resnet18_2d import build as build_dummy_2d
+
+
+_FIXTURE_MRI = Path(__file__).resolve().parents[1] / "fixtures" / "mri_sample" / "subject_0.nii.gz"
+
+
+class TestSelector:
+    def test_default_kind_is_volumetric(self, monkeypatch) -> None:
+        monkeypatch.delenv("MRI_MODEL_KIND", raising=False)
+        assert mri_selector.current_kind() == "volumetric_onnx"
+
+    def test_explicit_2d_selection(self, monkeypatch) -> None:
+        monkeypatch.setenv("MRI_MODEL_KIND", "resnet18_2d")
+        assert mri_selector.current_kind() == "resnet18_2d"
+
+    def test_unknown_kind_raises(self, monkeypatch) -> None:
+        monkeypatch.setenv("MRI_MODEL_KIND", "neural_net_supreme")
+        with pytest.raises(ValueError, match="unknown MRI_MODEL_KIND"):
+            mri_selector.current_kind()
+
+    def test_predict_routes_to_volumetric(self, monkeypatch, tmp_path) -> None:
+        monkeypatch.setenv("MRI_MODEL_KIND", "volumetric_onnx")
+        artifact = build_dummy_3d(tmp_path / "vol.onnx")
+        result = mri_selector.predict(
+            input_path=_FIXTURE_MRI,
+            checkpoint_path=artifact,
+            target_shape=(8, 8, 8),
+            label_names=("control", "abnormal"),
+        )
+        assert result["label_text"] in {"control", "abnormal"}
+
+    def test_predict_routes_to_2d(self, monkeypatch, tmp_path) -> None:
+        monkeypatch.setenv("MRI_MODEL_KIND", "resnet18_2d")
+        artifact = build_dummy_2d(tmp_path / "best.pt")
+        img_path = tmp_path / "scan.png"
+        Image.fromarray((np.random.RandomState(0).rand(160, 160, 3) * 255).astype("uint8")).save(str(img_path))
+        result = mri_selector.predict(
+            input_path=img_path,
+            checkpoint_path=artifact,
+        )
+        assert result["label_text"] in mri_selector.label_names_for_kind("resnet18_2d")