feat(models): add 2D resnet18 4-class Alzheimer's MRI inference module

src/models/mri_dl_2d.py

@@ -0,0 +1,86 @@
+"""Pretrained 2D MRI Alzheimer's classifier (resnet18, 4 classes).
+
+Decision-layer bridge for an externally-trained PyTorch checkpoint. Loads
+either a state_dict (default) or a full pickled model, applies the trainer's
+preprocessing (resize image_size=160, ImageNet normalisation), and emits the
+same dict shape as src.models.mri_model.predict_with_proba so downstream
+code paths don't care which backend produced the prediction.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+from torchvision import models, transforms
+
+from src.core.logger import get_logger
+
+logger = get_logger(__name__)
+
+CLASS_TO_IDX: dict[str, int] = {
+    "MildDemented": 0,
+    "ModerateDemented": 1,
+    "NonDemented": 2,
+    "VeryMildDemented": 3,
+}
+IDX_TO_CLASS: dict[int, str] = {v: k for k, v in CLASS_TO_IDX.items()}
+
+DEFAULT_IMAGE_SIZE = 160
+_IMAGENET_MEAN = (0.485, 0.456, 0.406)
+_IMAGENET_STD = (0.229, 0.224, 0.225)
+
+_TRANSFORM = transforms.Compose([
+    transforms.Resize((DEFAULT_IMAGE_SIZE, DEFAULT_IMAGE_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
+])
+
+
+def _build_resnet18_4class() -> nn.Module:
+    model = models.resnet18(weights=None)
+    model.fc = nn.Linear(model.fc.in_features, len(CLASS_TO_IDX))
+    return model
+
+
+def load(path: Path) -> nn.Module:
+    """Load checkpoint. Supports state_dict (preferred) or full pickled model."""
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"MRI 2D checkpoint not found: {path}")
+    obj = torch.load(str(path), map_location="cpu", weights_only=False)
+    if isinstance(obj, nn.Module):
+        model = obj
+    else:
+        model = _build_resnet18_4class()
+        clean = {k.removeprefix("module."): v for k, v in obj.items()}
+        model.load_state_dict(clean, strict=True)
+    model.eval()
+    return model
+
+
+def predict_image(model: nn.Module, image_path: Path) -> dict[str, Any]:
+    """Run inference on one image. Output mirrors mri_model.predict_with_proba."""
+    image_path = Path(image_path)
+    if not image_path.exists():
+        raise FileNotFoundError(f"MRI image not found: {image_path}")
+    img = Image.open(str(image_path)).convert("RGB")
+    tensor = _TRANSFORM(img).unsqueeze(0)
+
+    with torch.inference_mode():
+        logits = model(tensor)
+        probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy()
+
+    label_idx = int(np.argmax(probs))
+    return {
+        "label": label_idx,
+        "label_text": IDX_TO_CLASS[label_idx],
+        "confidence": float(probs[label_idx]),
+        "probabilities": [
+            {"label": i, "label_text": IDX_TO_CLASS[i], "probability": float(p)}
+            for i, p in enumerate(probs)
+        ],
+    }

tests/fixtures/build_dummy_resnet18_2d.py

@@ -0,0 +1,19 @@
+"""Build a randomly-initialised 4-class resnet18 state_dict for tests."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from torchvision import models
+
+
+def build(path: Path) -> Path:
+    """Save a state_dict at `path` and return the path. Idempotent."""
+    path = Path(path)
+    if path.exists():
+        return path
+    path.parent.mkdir(parents=True, exist_ok=True)
+    model = models.resnet18(weights=None)
+    model.fc = torch.nn.Linear(model.fc.in_features, 4)
+    torch.save(model.state_dict(), str(path))
+    return path

tests/models/test_mri_dl_2d.py

@@ -0,0 +1,61 @@
+"""Tests for src.models.mri_dl_2d — pretrained 4-class Alzheimer's resnet18."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from src.models import mri_dl_2d
+from tests.fixtures.build_dummy_resnet18_2d import build as build_dummy_2d
+
+
+def _png(path: Path, size: tuple[int, int] = (200, 200)) -> Path:
+    arr = (np.random.RandomState(0).rand(size[1], size[0], 3) * 255).astype(np.uint8)
+    Image.fromarray(arr, mode="RGB").save(str(path))
+    return path
+
+
+class TestMRIDL2D:
+    def test_class_to_idx_matches_trainer(self) -> None:
+        assert mri_dl_2d.CLASS_TO_IDX == {
+            "MildDemented": 0,
+            "ModerateDemented": 1,
+            "NonDemented": 2,
+            "VeryMildDemented": 3,
+        }
+
+    def test_idx_to_class_is_consistent(self) -> None:
+        for name, idx in mri_dl_2d.CLASS_TO_IDX.items():
+            assert mri_dl_2d.IDX_TO_CLASS[idx] == name
+
+    def test_load_missing_artifact_raises(self, tmp_path: Path) -> None:
+        with pytest.raises(FileNotFoundError, match="MRI 2D checkpoint not found"):
+            mri_dl_2d.load(tmp_path / "nope.pt")
+
+    def test_predict_image_returns_full_probs(self, tmp_path: Path) -> None:
+        ckpt = build_dummy_2d(tmp_path / "best.pt")
+        model = mri_dl_2d.load(ckpt)
+        img = _png(tmp_path / "scan.png")
+
+        result = mri_dl_2d.predict_image(model, img)
+
+        assert set(result) == {"label", "label_text", "confidence", "probabilities"}
+        assert result["label"] in {0, 1, 2, 3}
+        assert result["label_text"] in mri_dl_2d.CLASS_TO_IDX
+        assert 0.0 <= result["confidence"] <= 1.0
+        probs = result["probabilities"]
+        assert len(probs) == 4
+        assert abs(sum(p["probability"] for p in probs) - 1.0) < 1e-5
+        assert {p["label_text"] for p in probs} == set(mri_dl_2d.CLASS_TO_IDX)
+
+    def test_predict_works_for_grayscale_input(self, tmp_path: Path) -> None:
+        ckpt = build_dummy_2d(tmp_path / "best.pt")
+        model = mri_dl_2d.load(ckpt)
+        gray = (np.random.RandomState(1).rand(180, 180) * 255).astype(np.uint8)
+        path = tmp_path / "gray.png"
+        Image.fromarray(gray, mode="L").save(str(path))
+
+        result = mri_dl_2d.predict_image(model, path)
+        assert 0.0 <= result["confidence"] <= 1.0