feat(api): dispatch /predict/mri via MRI_MODEL_KIND env var

src/api/routes.py

@@ -319,25 +319,44 @@ def predict_bbb(req: BBBPredictRequest) -> BBBPredictResponse:
 
 @predict_router.post("/mri", response_model=MRIPredictResponse)
 def predict_mri(req: MRIPredictRequest) -> MRIPredictResponse:
-    """Predict from one MRI NIfTI image using an externally-trained ONNX model."""
-    artifact = _mri_model_path()
+    """Predict from one MRI image. Backend selected by MRI_MODEL_KIND env.
+
+    - `volumetric_onnx` (default): NIfTI volume + externally-trained ONNX.
+    - `resnet18_2d`: 2D image (.png/.jpg) + PyTorch state_dict, 4-class
+      Alzheimer's classifier (MildDemented/ModerateDemented/NonDemented/VeryMildDemented).
+    """
+    from src.models import mri_selector
+
+    kind = mri_selector.current_kind()
+    if kind == "resnet18_2d":
+        artifact = Path(os.environ.get(
+            "MRI_MODEL_PATH_2D", "data/processed/mri_dl_2d/best_model.pt",
+        ))
+    else:
+        artifact = _mri_model_path()
+
     if not artifact.exists():
         raise HTTPException(
             status_code=503,
             detail=(
-                f"MRI model artifact not available at {artifact}. "
-                "Export the trained volumetric model to ONNX and place it there, "
-                "or set MRI_MODEL_PATH."
+                f"MRI model artifact not available at {artifact} (kind={kind}). "
+                "Drop the trained checkpoint at this path, or override the path "
+                "via MRI_MODEL_PATH (3D ONNX) or MRI_MODEL_PATH_2D (2D resnet18)."
             ),
         )
     try:
-        model = mri_model.load(artifact)
-        pred = mri_model.predict_nifti(
-            model,
-            Path(req.input_path),
-            target_shape=req.target_shape,
-            label_names=req.label_names,
-        )
+        if kind == "resnet18_2d":
+            pred = mri_selector.predict(
+                input_path=Path(req.input_path),
+                checkpoint_path=artifact,
+            )
+        else:
+            pred = mri_selector.predict(
+                input_path=Path(req.input_path),
+                checkpoint_path=artifact,
+                target_shape=tuple(req.target_shape),
+                label_names=tuple(req.label_names) if req.label_names else None,
+            )
     except FileNotFoundError as e:
         raise HTTPException(status_code=404, detail=str(e))
     except ValueError as e:

src/api/schemas.py

@@ -115,7 +115,14 @@ class BBBPredictResponse(BaseModel):
 
 class MRIPredictRequest(BaseModel):
     """Single-subject MRI image prediction request."""
-    input_path: str = Field(..., description="Path to one .nii or .nii.gz MRI volume")
+    input_path: str = Field(
+        ...,
+        description=(
+            "Path to MRI input. With MRI_MODEL_KIND=volumetric_onnx (default), "
+            "expects a .nii/.nii.gz volume. With MRI_MODEL_KIND=resnet18_2d, "
+            "expects a 2D image (.png/.jpg)."
+        ),
+    )
     target_shape: tuple[int, int, int] = Field(
         (64, 64, 64),
         description="Model preprocessing resize target as (D, H, W)",

tests/api/test_mri_2d_route.py

@@ -0,0 +1,43 @@
+"""Integration: POST /predict/mri with MRI_MODEL_KIND=resnet18_2d."""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from fastapi.testclient import TestClient
+from PIL import Image
+
+from src.api.main import app
+from tests.fixtures.build_dummy_resnet18_2d import build as build_dummy_2d
+
+
+@pytest.fixture()
+def client_2d(monkeypatch, tmp_path):
+    monkeypatch.setenv("MRI_MODEL_KIND", "resnet18_2d")
+    ckpt = build_dummy_2d(tmp_path / "best.pt")
+    monkeypatch.setenv("MRI_MODEL_PATH_2D", str(ckpt))
+    return TestClient(app)
+
+
+def test_predict_mri_2d_happy_path(client_2d, tmp_path):
+    img_path = tmp_path / "scan.png"
+    Image.fromarray((np.random.RandomState(0).rand(170, 170, 3) * 255).astype("uint8")).save(str(img_path))
+
+    r = client_2d.post("/predict/mri", json={"input_path": str(img_path)})
+    assert r.status_code == 200, r.text
+    data = r.json()
+    assert data["label_text"] in {
+        "MildDemented", "ModerateDemented", "NonDemented", "VeryMildDemented",
+    }
+    assert 0.0 <= data["confidence"] <= 1.0
+    assert len(data["probabilities"]) == 4
+
+
+def test_predict_mri_2d_missing_artifact_returns_503(monkeypatch, tmp_path):
+    monkeypatch.setenv("MRI_MODEL_KIND", "resnet18_2d")
+    monkeypatch.setenv("MRI_MODEL_PATH_2D", str(tmp_path / "missing.pt"))
+    img_path = tmp_path / "scan.png"
+    Image.fromarray((np.random.RandomState(0).rand(170, 170, 3) * 255).astype("uint8")).save(str(img_path))
+    client = TestClient(app)
+    r = client.post("/predict/mri", json={"input_path": str(img_path)})
+    assert r.status_code == 503
+    assert "kind=resnet18_2d" in r.text