Fix user-facing API bugs

brats/__init__.py

@@ -13,4 +13,4 @@
     PediatricSegmenter,
 )
 
-logger.remove()
+logger.disable("brats")

brats/core/brats_algorithm.py

@@ -97,16 +97,21 @@ def _process_single_output(
         if self.task == Task.MISSING_MRI:
             # Missing MRI has no fixed names since the missing modality
             # differs and is included in the name
-            algorithm_output = Path(tmp_output_folder).iterdir().__next__()
+            algorithm_output = next(Path(tmp_output_folder).iterdir(), None)
+            if algorithm_output is None:
+                raise FileNotFoundError(
+                    f"No output found for subject {subject_id} in {tmp_output_folder}"
+                )
         else:
             # extract id from subject id, i.e. BraTS-MEN-00000-000 => 00000-000
             identifier = self.extract_identifier_from_subject_id(subject_id)
-            possible_output = list(Path(tmp_output_folder).glob(f"*{identifier}*"))
-            if len(possible_output) == 0:
+            algorithm_output = next(
+                Path(tmp_output_folder).glob(f"*{identifier}*"), None
+            )
+            if algorithm_output is None:
                 raise FileNotFoundError(
                     f"No output found for subject {subject_id} in {tmp_output_folder}"
                 )
-            algorithm_output = possible_output[0]
 
         # ensure path exists and rename output to the desired path
         output_file = Path(output_file).absolute()
@@ -134,9 +139,14 @@ def _process_batch_output(
             if self.task == Task.MISSING_MRI:
                 # Missing MRI has no fixed names since the missing modality differs
                 # and is included in the name
-                algorithm_output = (
-                    Path(tmp_output_folder).glob(f"*{internal_name}*").__next__()
+                algorithm_output = next(
+                    Path(tmp_output_folder).glob(f"*{internal_name}*"), None
                 )
+                if algorithm_output is None:
+                    raise FileNotFoundError(
+                        f"No output found for subject {internal_name} "
+                        f"in {tmp_output_folder}"
+                    )
                 try:
                     modality = algorithm_output.name.split("-")[-1].split(".")[0]
                 except IndexError:
@@ -150,13 +160,14 @@ def _process_batch_output(
                 )
             else:
                 identifier = self.extract_identifier_from_subject_id(internal_name)
-                possible_outputs = list(Path(tmp_output_folder).glob(f"*{identifier}*"))
-                if len(possible_outputs) == 0:
+                algorithm_output = next(
+                    Path(tmp_output_folder).glob(f"*{identifier}*"), None
+                )
+                if algorithm_output is None:
                     logger.error(
                         f"No output found for subject {internal_name} in {tmp_output_folder}"
                     )
                     continue
-                algorithm_output = possible_outputs[0]
 
                 output_file = output_folder / f"{external_name}.nii.gz"
             shutil.move(algorithm_output, output_file)

brats/preprocessing.py

@@ -5,8 +5,14 @@
 
 from brats.constants import (
     AdultGliomaPreAndPostTreatmentAlgorithms,
+    AdultGliomaPreTreatmentAlgorithms,
+    AfricaAlgorithms,
     Algorithms,
+    GoATAlgorithms,
+    InpaintingAlgorithms,
+    MeningiomaAlgorithms,
     MeningiomaRTAlgorithms,
+    MetastasesAlgorithms,
     MissingMRIAlgorithms,
     PediatricAlgorithms,
 )
@@ -373,6 +379,7 @@ def preprocess_for_challenge(
 
     Raises:
         ValueError: If required modalities are missing for the challenge
+        TypeError: If challenge is not a supported Algorithms enum member
     """
     challenge_name = str(challenge)
 
@@ -392,24 +399,24 @@ def _require_all_modalities() -> list[str | Path]:
             raise ValueError(
                 f"All modalities required for {challenge_name} preprocessing"
             )
-        return cast(list[str | Path], all_paths)
+        return cast(list[Union[str, Path]], all_paths)
 
     # Route to appropriate preprocessing function
-    if str(AdultGliomaPreAndPostTreatmentAlgorithms.__name__) in challenge_name:
+    if isinstance(challenge, AdultGliomaPreAndPostTreatmentAlgorithms):
         paths = _require_all_modalities()
         preprocess_coreg_mni152reg_bet(
             *paths,
             normalizer=normalizer,
         )
 
-    elif str(PediatricAlgorithms.__name__) in challenge_name:
+    elif isinstance(challenge, PediatricAlgorithms):
         paths = _require_all_modalities()
         preprocess_coreg_sri24reg_defacing(
             *paths,
             normalizer=normalizer,
         )
 
-    elif str(MissingMRIAlgorithms.__name__) in challenge_name:
+    elif isinstance(challenge, MissingMRIAlgorithms):
         preprocess_coreg_sri24reg_bet_allow_missing(
             t1_input,
             t1c_input,
@@ -421,7 +428,7 @@ def _require_all_modalities() -> list[str | Path]:
             flair_output,
             normalizer=normalizer,
         )
-    elif str(MeningiomaRTAlgorithms.__name__) in challenge_name:
+    elif isinstance(challenge, MeningiomaRTAlgorithms):
         if t1c_input is None or t1c_output is None:
             raise ValueError(
                 f"T1c modality required for {challenge_name} preprocessing"
@@ -431,9 +438,24 @@ def _require_all_modalities() -> list[str | Path]:
             t1c_output=t1c_output,
             normalizer=normalizer,
         )
-    else:  # Most challenges use SRI24 with BET
+    elif isinstance(
+        challenge,
+        (
+            AdultGliomaPreTreatmentAlgorithms,
+            AfricaAlgorithms,
+            GoATAlgorithms,
+            InpaintingAlgorithms,
+            MeningiomaAlgorithms,
+            MetastasesAlgorithms,
+        ),
+    ):
         paths = _require_all_modalities()
         preprocess_coreg_sri24reg_bet(
             *paths,
             normalizer=normalizer,
         )
+    else:
+        raise TypeError(
+            "challenge must be a supported Algorithms enum member, "
+            f"not {type(challenge).__name__}"
+        )

docs/tutorials/tutorial.ipynb

@@ -552,9 +552,10 @@
     }
    ],
    "source": [
-    "from brats.utils.logging import add_console_handler, remove_console_handler\n",
+     "from brats.utils.logging import add_console_handler, enable, remove_console_handler\n",
     "\n",
-    "add_console_handler(level=\"INFO\")  # Set the desired log level\n",
+     "enable()\n",
+     "add_console_handler(level=\"INFO\")  # Set the desired log level\n",
     "segmenter = AdultGliomaPreTreatmentSegmenter(\n",
     "    algorithm=AdultGliomaPreTreatmentAlgorithms.BraTS23_3,  # Use the 3rd placed algorithm of the Adult Glioma BraTS 2023 challenge\n",
     "    cuda_devices=\"1\",  # Select GPU device with ID 1\n",

tests/core/test_missing_mri_algorithms.py

@@ -94,6 +94,52 @@ def test_standardize_segmentation_inputs_list(self, mock_standardize_single_inpu
         )
         mock_standardize_single_inputs.assert_called_once()
 
+    @patch("brats.core.missing_mri_algorithms.input_sanity_check")
+    @patch("brats.core.brats_algorithm.run_docker_container")
+    @patch("brats.core.brats_algorithm.InferenceSetup")
+    def test_infer_single_raises_file_not_found_for_missing_output(
+        self, mock_inference_setup, mock_run_container, mock_input_sanity_check
+    ):
+        output_folder = self.test_dir / "single_tmp_output"
+        output_folder.mkdir()
+        mock_inference_setup.return_value.__enter__.return_value = (
+            self.tmp_data_folder,
+            output_folder,
+        )
+
+        with self.assertRaisesRegex(FileNotFoundError, "No output found for subject"):
+            self.missing_mri.infer_single(
+                t1c=self.t1c,
+                t1n=self.t1n,
+                t2w=self.t2w,
+                output_file=self.test_dir / "output.nii.gz",
+            )
+
+        mock_run_container.assert_called_once()
+        mock_input_sanity_check.assert_called_once()
+
+    @patch("brats.core.missing_mri_algorithms.input_sanity_check")
+    @patch("brats.core.brats_algorithm.run_docker_container")
+    @patch("brats.core.brats_algorithm.InferenceSetup")
+    def test_infer_batch_raises_file_not_found_for_missing_output(
+        self, mock_inference_setup, mock_run_container, mock_input_sanity_check
+    ):
+        output_folder = self.test_dir / "batch_tmp_output"
+        output_folder.mkdir()
+        mock_inference_setup.return_value.__enter__.return_value = (
+            self.tmp_data_folder,
+            output_folder,
+        )
+
+        with self.assertRaisesRegex(FileNotFoundError, "No output found for subject"):
+            self.missing_mri.infer_batch(
+                data_folder=self.data_folder,
+                output_folder=self.test_dir / "outputs",
+            )
+
+        mock_run_container.assert_called_once()
+        mock_input_sanity_check.assert_called_once()
+
     # Initialization tests
 
     def test_missing_mri_initialization(self):

tests/test_preprocessing.py

@@ -0,0 +1,151 @@
+import importlib
+import sys
+from contextlib import ExitStack
+from types import ModuleType
+from unittest.mock import patch
+
+import pytest
+
+from brats.constants import (
+    AdultGliomaPreAndPostTreatmentAlgorithms,
+    AdultGliomaPreTreatmentAlgorithms,
+    AfricaAlgorithms,
+    GoATAlgorithms,
+    InpaintingAlgorithms,
+    MeningiomaAlgorithms,
+    MeningiomaRTAlgorithms,
+    MetastasesAlgorithms,
+    MissingMRIAlgorithms,
+    PediatricAlgorithms,
+)
+
+
+@pytest.fixture
+def preprocessing_module(monkeypatch):
+    """Import preprocessing with a lightweight fake optional dependency."""
+    brainles_preprocessing = ModuleType("brainles_preprocessing")
+    constants = ModuleType("brainles_preprocessing.constants")
+    modality = ModuleType("brainles_preprocessing.modality")
+    normalization = ModuleType("brainles_preprocessing.normalization")
+    preprocessor = ModuleType("brainles_preprocessing.preprocessor")
+
+    class Atlas:
+        BRATS_MNI152 = "MNI152"
+        BRATS_SRI24 = "SRI24"
+
+    class CenterModality:
+        pass
+
+    class Modality:
+        pass
+
+    class Normalizer:
+        pass
+
+    class AtlasCentricPreprocessor:
+        pass
+
+    class NativeSpacePreprocessor:
+        pass
+
+    constants.Atlas = Atlas
+    modality.CenterModality = CenterModality
+    modality.Modality = Modality
+    normalization.Normalizer = Normalizer
+    preprocessor.AtlasCentricPreprocessor = AtlasCentricPreprocessor
+    preprocessor.NativeSpacePreprocessor = NativeSpacePreprocessor
+
+    for module in [
+        brainles_preprocessing,
+        constants,
+        modality,
+        normalization,
+        preprocessor,
+    ]:
+        monkeypatch.setitem(sys.modules, module.__name__, module)
+
+    sys.modules.pop("brats.preprocessing", None)
+    module = importlib.import_module("brats.preprocessing")
+    yield module
+    sys.modules.pop("brats.preprocessing", None)
+
+
+def _all_modalities() -> dict[str, str]:
+    return {
+        "t1_input": "t1.nii.gz",
+        "t1c_input": "t1c.nii.gz",
+        "t2_input": "t2.nii.gz",
+        "flair_input": "flair.nii.gz",
+        "t1_output": "t1-out.nii.gz",
+        "t1c_output": "t1c-out.nii.gz",
+        "t2_output": "t2-out.nii.gz",
+        "flair_output": "flair-out.nii.gz",
+    }
+
+
+@pytest.mark.parametrize(
+    ("challenge", "expected_pipeline"),
+    [
+        (
+            AdultGliomaPreAndPostTreatmentAlgorithms.BraTS25_1,
+            "preprocess_coreg_mni152reg_bet",
+        ),
+        (PediatricAlgorithms.BraTS25_1, "preprocess_coreg_sri24reg_defacing"),
+        (
+            MissingMRIAlgorithms.BraTS25_1,
+            "preprocess_coreg_sri24reg_bet_allow_missing",
+        ),
+        (MeningiomaRTAlgorithms.BraTS25_1, "preprocess_deface_only"),
+        (
+            AdultGliomaPreTreatmentAlgorithms.BraTS23_1,
+            "preprocess_coreg_sri24reg_bet",
+        ),
+        (AfricaAlgorithms.BraTS25_1, "preprocess_coreg_sri24reg_bet"),
+        (GoATAlgorithms.BraTS25_1A, "preprocess_coreg_sri24reg_bet"),
+        (InpaintingAlgorithms.BraTS25_1A, "preprocess_coreg_sri24reg_bet"),
+        (MeningiomaAlgorithms.BraTS25_1, "preprocess_coreg_sri24reg_bet"),
+        (MetastasesAlgorithms.BraTS23_1, "preprocess_coreg_sri24reg_bet"),
+    ],
+)
+def test_preprocess_for_challenge_dispatches_by_enum(
+    preprocessing_module, challenge, expected_pipeline
+):
+    pipeline_names = [
+        "preprocess_coreg_mni152reg_bet",
+        "preprocess_coreg_sri24reg_defacing",
+        "preprocess_coreg_sri24reg_bet_allow_missing",
+        "preprocess_deface_only",
+        "preprocess_coreg_sri24reg_bet",
+    ]
+    paths = _all_modalities()
+
+    with ExitStack() as stack:
+        pipelines = {
+            name: stack.enter_context(patch.object(preprocessing_module, name))
+            for name in pipeline_names
+        }
+        preprocessing_module.preprocess_for_challenge(challenge, **paths)
+
+    pipelines[expected_pipeline].assert_called_once()
+    for name, pipeline in pipelines.items():
+        if name != expected_pipeline:
+            pipeline.assert_not_called()
+
+
+def test_preprocess_for_challenge_rejects_raw_value(preprocessing_module):
+    with pytest.raises(TypeError, match="supported Algorithms enum member"):
+        preprocessing_module.preprocess_for_challenge(
+            AdultGliomaPreAndPostTreatmentAlgorithms.BraTS25_1.value,
+            **_all_modalities(),
+        )
+
+
+def test_preprocess_for_challenge_validates_required_modalities(preprocessing_module):
+    paths = _all_modalities()
+    paths["t1_input"] = None
+
+    with pytest.raises(ValueError, match="All modalities required"):
+        preprocessing_module.preprocess_for_challenge(
+            AdultGliomaPreAndPostTreatmentAlgorithms.BraTS25_1,
+            **paths,
+        )