diff --git a/brats/__init__.py b/brats/__init__.py index a67c496..dba2e29 100644 --- a/brats/__init__.py +++ b/brats/__init__.py @@ -13,4 +13,4 @@ PediatricSegmenter, ) -logger.remove() +logger.disable("brats") diff --git a/brats/core/brats_algorithm.py b/brats/core/brats_algorithm.py index 6490678..4a36e0b 100644 --- a/brats/core/brats_algorithm.py +++ b/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) diff --git a/brats/preprocessing.py b/brats/preprocessing.py index 574b80a..a7d9fcb 100644 --- a/brats/preprocessing.py +++ b/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__}" + ) diff --git a/docs/tutorials/tutorial.ipynb b/docs/tutorials/tutorial.ipynb index 3ec9ac6..313640f 100644 --- a/docs/tutorials/tutorial.ipynb +++ b/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", diff --git a/tests/core/test_missing_mri_algorithms.py b/tests/core/test_missing_mri_algorithms.py index e93d7c5..b3ff4a0 100644 --- a/tests/core/test_missing_mri_algorithms.py +++ b/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): diff --git a/tests/test_preprocessing.py b/tests/test_preprocessing.py new file mode 100644 index 0000000..9ff12b1 --- /dev/null +++ b/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, + ) diff --git a/tests/utils/test_data_handling.py b/tests/utils/test_data_handling.py index 3a08101..28a2fc2 100644 --- a/tests/utils/test_data_handling.py +++ b/tests/utils/test_data_handling.py @@ -13,6 +13,7 @@ input_sanity_check, remove_tmp_folder, ) +from brats.utils.logging import disable, enable class TestDataHandlingUtils(unittest.TestCase): @@ -93,12 +94,16 @@ def test_inference_setup_error_preserves_temp_log(self): def test_inference_setup_with_log_file_on_error(self): tmp_log_file = self.test_dir / "error.log" - with self.assertRaises(RuntimeError): # noqa: SIM117 - with InferenceSetup(log_file=tmp_log_file) as ( - _tmp_data_folder, - _tmp_output_folder, - ): - raise RuntimeError("test error") + enable() + try: + with self.assertRaises(RuntimeError): # noqa: SIM117 + with InferenceSetup(log_file=tmp_log_file) as ( + _tmp_data_folder, + _tmp_output_folder, + ): + raise RuntimeError("test error") + finally: + disable() self.assertTrue(tmp_log_file.exists()) self.assertGreater(tmp_log_file.stat().st_size, 0) tmp_log_file.unlink() diff --git a/tests/utils/test_logging.py b/tests/utils/test_logging.py index 850f0c7..bb74791 100644 --- a/tests/utils/test_logging.py +++ b/tests/utils/test_logging.py @@ -1,6 +1,10 @@ +import importlib +from io import StringIO + import pytest from loguru import logger +import brats from brats.utils.logging import ( _reset_logging_state_for_tests, add_console_handler, @@ -37,6 +41,18 @@ def fake_enable(name): assert enabled_modules.get("brats") is True +def test_import_brats_preserves_existing_loguru_handler(): + sink = StringIO() + handler_id = logger.add(sink, format="{message}") + + try: + importlib.reload(brats) + logger.info("handler survived import") + assert "handler survived import" in sink.getvalue() + finally: + logger.remove(handler_id) + + def test_add_console_handler_writes_to_stderr(capfd): add_console_handler(level="INFO")