feat(mri): add mask_brain (intensity threshold + morphological opening)

src/pipelines/mri_pipeline.py

@@ -13,8 +13,10 @@ from __future__ import annotations
 
 import os
 
+import nibabel as nib
 import numpy as np
 import pyarrow as pa
+from scipy import ndimage as scipy_ndimage
 
 from src.core.logger import get_logger
 
@@ -48,3 +50,32 @@ def is_valid_volume(volume: np.ndarray | None) -> bool:
     if not np.all(np.isfinite(volume)):
         return False
     return True
+
+
+def mask_brain(
+    volume: np.ndarray,
+    intensity_threshold: float | None = None,
+) -> np.ndarray:
+    """Build a brain mask from a 3-D MRI volume.
+
+    Two-step pipeline:
+      1. Intensity threshold: keep voxels above `intensity_threshold`. When
+         `None`, use the volume's mean as a robust auto-threshold (works on
+         the synthetic fixture where brain ≫ background; for real data the
+         caller should pass an Otsu or BET-derived threshold explicitly).
+      2. Morphological opening (`scipy.ndimage.binary_opening`) to remove
+         isolated noise voxels and disconnected fragments.
+
+    Args:
+        volume: 3-D numeric `np.ndarray` (must satisfy `is_valid_volume`).
+        intensity_threshold: Voxel-intensity floor. `None` → use `volume.mean()`.
+
+    Returns:
+        A boolean `np.ndarray` of the same shape as `volume`. True = brain.
+    """
+    if intensity_threshold is None:
+        intensity_threshold = float(volume.mean())
+
+    raw = volume > intensity_threshold
+    cleaned = scipy_ndimage.binary_opening(raw, iterations=1)
+    return cleaned.astype(bool)

tests/pipelines/test_mri_pipeline.py

@@ -3,10 +3,14 @@ from __future__ import annotations
 
 from pathlib import Path
 
+import nibabel as nib
 import numpy as np
 import pytest
 
-from src.pipelines.mri_pipeline import is_valid_volume
+from src.pipelines.mri_pipeline import (
+    is_valid_volume,
+    mask_brain,
+)
 
 
 FIXTURE_DIR = Path(__file__).parent.parent / "fixtures" / "mri_sample"
@@ -45,3 +49,57 @@ class TestIsValidVolume:
     def test_rejects_non_array(self) -> None:
         assert is_valid_volume([[[1, 2]], [[3, 4]]]) is False
         assert is_valid_volume(None) is False
+
+
+class TestMaskBrain:
+    def _load_subject(self, sid: str) -> np.ndarray:
+        return nib.load(FIXTURE_DIR / f"{sid}.nii.gz").get_fdata()
+
+    def test_returns_bool_mask_of_same_shape(self) -> None:
+        vol = self._load_subject("subject_0")
+        mask = mask_brain(vol)
+        assert isinstance(mask, np.ndarray)
+        assert mask.dtype == bool
+        assert mask.shape == vol.shape
+
+    def test_mask_separates_brain_from_background(self) -> None:
+        """Default threshold should keep the spherical-brain center voxels in."""
+        vol = self._load_subject("subject_0")
+        mask = mask_brain(vol)
+        # The fixture's brain region (radius 3 around center) intensity is ~10;
+        # background is ~0.1. Some brain voxels MUST survive the mask.
+        assert mask.sum() > 0
+        # The center voxel (always brain) MUST be in the mask.
+        center = tuple(s // 2 for s in vol.shape)
+        assert mask[center]
+
+    def test_mask_drops_low_intensity_background(self) -> None:
+        """Voxels with intensity well below the brain core must be excluded."""
+        vol = self._load_subject("subject_0")
+        mask = mask_brain(vol, intensity_threshold=5.0)
+        # Background voxels (intensity ~0.1) must NOT be in the mask.
+        bg_voxel = (0, 0, 0)
+        assert mask[bg_voxel] == False  # noqa: E712
+
+    def test_explicit_threshold_overrides_default(self) -> None:
+        vol = self._load_subject("subject_0")
+        # A very high threshold should produce far fewer mask voxels.
+        mask_default = mask_brain(vol)
+        mask_strict = mask_brain(vol, intensity_threshold=100.0)
+        assert mask_strict.sum() < mask_default.sum()
+
+    def test_does_not_mutate_input(self) -> None:
+        vol = self._load_subject("subject_0")
+        original = vol.copy()
+        _ = mask_brain(vol)
+        np.testing.assert_array_equal(vol, original)
+
+    def test_morphological_cleanup_removes_isolated_voxels(self) -> None:
+        """A single bright voxel surrounded by background must be removed by the
+        opening-style morphological cleanup."""
+        vol = np.zeros((8, 8, 8), dtype=np.float64)
+        vol[4, 4, 4] = 100.0
+        mask = mask_brain(vol, intensity_threshold=50.0)
+        # Without cleanup, the single voxel would survive. With morphological
+        # opening, it must be removed.
+        assert mask.sum() == 0