BraTS 0.0.0.dev0__py3-none-any.whl

brats/__init__.py

@@ -0,0 +1,16 @@
+from loguru import logger
+
+from brats.core.inpainting_algorithms import Inpainter
+from brats.core.missing_mri_algorithms import MissingMRI
+from brats.core.segmentation_algorithms import (
+    AdultGliomaPreTreatmentSegmenter,
+    AdultGliomaPreAndPostTreatmentSegmenter,
+    AfricaSegmenter,
+    GoATSegmenter,
+    MeningiomaSegmenter,
+    MetastasesSegmenter,
+    MeningiomaRTSegmenter,
+    PediatricSegmenter,
+)
+
+logger.remove()

brats/constants.py

@@ -0,0 +1,268 @@
+from enum import Enum
+from pathlib import Path
+
+# TASK ENUM
+
+
+class Backends(str, Enum):
+    """Available backends for running the algorithms."""
+
+    DOCKER = "docker"
+    """Run the algorithms using Docker containers."""
+
+    SINGULARITY = "singularity"
+    """Run the algorithms using Singularity containers."""
+
+
+class Task(str, Enum):
+    """Available tasks."""
+
+    SEGMENTATION = "SEGMENTATION"
+    """Segmentation task."""
+
+    INPAINTING = "INPAINTING"
+    """Inpainting task."""
+
+    MISSING_MRI = "MISSING_MRI"
+    """Missing MRI task."""
+
+
+# ALGORITHM ENUMS
+
+
+class Algorithms(str, Enum):
+    """Parent class for constants of the available algorithms."""
+
+    pass
+
+
+class AdultGliomaPreAndPostTreatmentAlgorithms(Algorithms):
+    """Constants for the available adult glioma pre and post treatment segmentation algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """ BraTS25 Adult Glioma Segmentation 1st place """
+    BraTS25_2 = "BraTS25_2"
+    """ BraTS25 Adult Glioma Segmentation 2nd place """
+    BraTS25_3A = "BraTS25_3A"
+    """ BraTS25 Adult Glioma Segmentation 3rd place (tie) """
+    BraTS25_3B = "BraTS25_3B"
+    """ BraTS25 Adult Glioma Segmentation 3rd place (tie) """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24 Adult Glioma Segmentation 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24 Adult Glioma Segmentation 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24 Adult Glioma Segmentation 3rd place """
+
+
+class AdultGliomaPreTreatmentAlgorithms(Algorithms):
+    """Constants for the available adult glioma pre treatment segmentation algorithms."""
+
+    BraTS23_1 = "BraTS23_1"
+    """BraTS23 Adult Glioma Segmentation 1st place (GPU only)"""
+    BraTS23_2 = "BraTS23_2"
+    """BraTS23 Adult Glioma Segmentation 2nd place (GPU only)"""
+    BraTS23_3 = "BraTS23_3"
+    """BraTS23 Adult Glioma Segmentation 3rd place (GPU only)"""
+
+
+class MeningiomaRTAlgorithms(Algorithms):
+    """Constants for the available meningioma segmentation - Radio Therapy algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """ BraTS25 Meningioma Segmentation 1st place """
+    BraTS25_2 = "BraTS25_2"
+    """ BraTS25 Meningioma Segmentation 2nd place """
+    BraTS25_3 = "BraTS25_3"
+    """ BraTS25 Meningioma Segmentation 3rd place """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24 Meningioma Segmentation 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24 Meningioma Segmentation 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24 Meningioma Segmentation 3rd place """
+
+
+class MeningiomaAlgorithms(Algorithms):
+    """Constants for the available meningioma segmentation algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """BraTS25 Meningioma Segmentation 1st place """
+    BraTS25_2 = "BraTS25_2"
+    """BraTS25 Meningioma Segmentation 2nd place """
+
+    BraTS23_1 = "BraTS23_1"
+    """BraTS23 Meningioma Segmentation 1st place (GPU only)"""
+    BraTS23_2 = "BraTS23_2"
+    """BraTS23 Meningioma Segmentation 2nd place (GPU and CPU)"""
+    BraTS23_3 = "BraTS23_3"
+    """BraTS23 Meningioma Segmentation 3rd place (GPU only)"""
+
+
+class PediatricAlgorithms(Algorithms):
+    """Constants for the available pediatric segmentation algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """ BraTS25 Pediatric Segmentation 1st place"""
+    BraTS25_2 = "BraTS25_2"
+    """ BraTS25 Pediatric Segmentation 2nd place"""
+    BraTS25_3 = "BraTS25_3"
+    """ BraTS25 Pediatric Segmentation 3rd place"""
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24 Pediatric Segmentation 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24 Pediatric Segmentation 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24 Pediatric Segmentation 3rd place """
+
+    BraTS23_1 = "BraTS23_1"
+    """BraTS23 Pediatric Segmentation 1st place (GPU only)"""
+    BraTS23_2 = "BraTS23_2"
+    """BraTS23 Pediatric Segmentation 2nd place (GPU only)"""
+    BraTS23_3 = "BraTS23_3"
+    """BraTS23 Pediatric Segmentation 3rd place (GPU only)"""
+
+
+class AfricaAlgorithms(Algorithms):
+    """Constants for the available africa segmentation algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """ BraTS25 BraTS-SSA Segmentation 1st place """
+    BraTS25_2A = "BraTS25_2A"
+    """ BraTS25 BraTS-SSA Segmentation 2nd place (Tie)"""
+    BraTS25_2B = "BraTS25_2B"
+    """ BraTS25 BraTS-SSA Segmentation 2nd place (Tie) """
+    BraTS25_3 = "BraTS25_3"
+    """ BraTS25 BraTS-SSA Segmentation 3rd place """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24 BraTS-Africa Segmentation 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24 BraTS-Africa Segmentation 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24 BraTS-Africa Segmentation 3rd place """
+
+    BraTS23_1 = "BraTS23_1"
+    """BraTS23 BraTS-Africa Segmentation 1st place (GPU only)"""
+    BraTS23_2 = "BraTS23_2"
+    """BraTS23 BraTS-Africa Segmentation 2nd place (GPU only)"""
+    BraTS23_3 = "BraTS23_3"
+    """BraTS23 BraTS-Africa Segmentation 3rd place (GPU and CPU)"""
+
+
+class MetastasesAlgorithms(Algorithms):
+    """Constants for the available Inpainting algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """BraTS25  Brain Metastases Segmentation 1st place"""
+    BraTS25_2 = "BraTS25_2"
+    """BraTS25  Brain Metastases Segmentation 2nd place"""
+
+    BraTS23_1 = "BraTS23_1"
+    """BraTS23  Brain Metastases Segmentation 1st place (GPU only)"""
+    BraTS23_2 = "BraTS23_2"
+    """BraTS23  Brain Metastases Segmentation 2nd place (GPU only)"""
+    BraTS23_3 = "BraTS23_3"
+    """BraTS23  Brain Metastases Segmentation 3rd place (GPU only)"""
+
+
+class InpaintingAlgorithms(Algorithms):
+    """Constants for the available BraTS Inpainting algorithms."""
+
+    BraTS25_1A = "BraTS25_1A"
+    """ BraTS25  Inpainting shared 1st place (tie) """
+    BraTS25_1B = "BraTS25_1B"
+    """ BraTS25  Inpainting shared 1st place (tie)"""
+    BraTS25_2 = "BraTS25_2"
+    """ BraTS25  Inpainting 2nd place """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24  Inpainting 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24  Inpainting 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24  Inpainting 3rd place """
+
+    BraTS23_1 = "BraTS23_1"
+    """ BraTS23  Inpainting 1st place """
+    BraTS23_2 = "BraTS23_2"
+    """ BraTS23  Inpainting 2nd place (Very Slow) """
+    BraTS23_3 = "BraTS23_3"
+    """ BraTS23  Inpainting 3rd place """
+
+
+class MissingMRIAlgorithms(Algorithms):
+    """Constants for the available missing mri algorithms."""
+
+    BraTS25_1 = "BraTS25_1"
+    """ BraTS25  MissingMRI 1st place """
+    BraTS25_2 = "BraTS25_2"
+    """ BraTS25  MissingMRI 2nd place """
+    BraTS25_3 = "BraTS25_3"
+    """ BraTS25  MissingMRI 3rd place """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24  MissingMRI 1st place """
+    BraTS24_2 = "BraTS24_2"
+    """ BraTS24  MissingMRI 2nd place """
+    BraTS24_3 = "BraTS24_3"
+    """ BraTS24  MissingMRI 3rd place """
+
+    BraTS23_1 = "BraTS23_1"
+    """ BraTS23  MissingMRI 1st place """
+
+
+class GoATAlgorithms(Algorithms):
+    """Constants for the available missing mri  algorithms."""
+
+    BraTS25_1A = "BraTS25_1A"
+    """ BraTS25 Generalizability Across Tumors (BraTS-GoAT) 1st place (tie) """
+    BraTS25_1B = "BraTS25_1B"
+    """ BraTS25 Generalizability Across Tumors (BraTS-GoAT) 1st place (tie) """
+    BraTS25_1C = "BraTS25_1C"
+    """ BraTS25 Generalizability Across Tumors (BraTS-GoAT) 1st place (tie) """
+    BraTS25_1D = "BraTS25_1D"
+    """ BraTS25 Generalizability Across Tumors (BraTS-GoAT) 1st place (tie) """
+
+    BraTS24_1 = "BraTS24_1"
+    """ BraTS24 Generalizability Across Tumors (BraTS-GoAT) 1st place (The only submission)"""
+
+
+# DIRECTORIES
+DATA_DIR = Path(__file__).parent / "data"
+META_DIR = DATA_DIR / "meta"
+PARAMETERS_DIR = DATA_DIR / "parameters"
+ADDITIONAL_FILES_FOLDER = DATA_DIR / "additional_files"
+
+# META DATA FILE PATHS
+ADULT_GLIOMA_PRE_TREATMENT_SEGMENTATION_ALGORITHMS = (
+    META_DIR / "adult_glioma_pre_treatment.yml"
+)
+ADULT_GLIOMA_PRE_AND_POST_TREATMENT_SEGMENTATION_ALGORITHMS = (
+    META_DIR / "adult_glioma_pre_and_post_treatment.yml"
+)
+MENINGIOMA_SEGMENTATION_ALGORITHMS = META_DIR / "meningioma.yml"
+MENINGIOMA_RT_SEGMENTATION_ALGORITHMS = META_DIR / "meningioma_rt.yml"
+PEDIATRIC_SEGMENTATION_ALGORITHMS = META_DIR / "pediatric.yml"
+AFRICA_SEGMENTATION_ALGORITHMS = META_DIR / "africa.yml"
+METASTASES_SEGMENTATION_ALGORITHMS = META_DIR / "metastases.yml"
+GOAT_SEGMENTATION_ALGORITHMS = META_DIR / "goat.yml"
+INPAINTING_ALGORITHMS = META_DIR / "inpainting.yml"
+MISSING_MRI_ALGORITHMS = META_DIR / "missing_mri.yml"
+
+DUMMY_PARAMETERS = PARAMETERS_DIR / "dummy.yml"
+
+# ZENODO
+ZENODO_RECORD_BASE_URL = "https://zenodo.org/api/records"
+
+OUTPUT_NAME_SCHEMA = {
+    Task.SEGMENTATION: "{subject_id}.nii.gz",
+    Task.INPAINTING: "{subject_id}-t1n-inference.nii.gz",
+}
+
+
+# PACKAGE citation
+PACKAGE_CITATION = "https://arxiv.org/abs/2506.13807"

brats/core/brats_algorithm.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import shutil
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, Optional, Callable
+
+from loguru import logger
+
+from brats.core.docker import run_container as run_docker_container
+from brats.core.singularity import run_container as run_singularity_container
+from brats.utils.algorithm_config import load_algorithms
+from brats.constants import OUTPUT_NAME_SCHEMA, Algorithms, Task, Backends
+from brats.utils.data_handling import InferenceSetup
+from brats.utils.exceptions import AlgorithmConfigException
+
+
+class BraTSAlgorithm(ABC):
+    """
+    This class serves as the basis for all BraTS algorithms. It provides a common interface and implements the logic for single and batch inference.
+    """
+
+    def __init__(
+        self,
+        algorithm: Algorithms,
+        algorithms_file_path: Path,
+        task: Task,
+        cuda_devices: str = "0",
+        force_cpu: bool = False,
+    ):
+        # inference device setup
+        self.force_cpu = force_cpu
+        self.cuda_devices = cuda_devices
+
+        self.task = task
+        self.algorithm_list = load_algorithms(file_path=algorithms_file_path)
+        # save algorithm identifier for logging etc.
+        self.algorithm_key = algorithm.value
+        if self.algorithm_key not in self.algorithm_list:
+            raise AlgorithmConfigException(
+                f"Algorithm {self.algorithm_key} not found in {algorithms_file_path}"
+            )
+        # data for selected algorithm
+        self.algorithm = self.algorithm_list[algorithm.value]
+
+        logger.info(
+            f"Instantiated {self.__class__.__name__} with algorithm: {self.algorithm_key} by {self.algorithm.meta.authors}"
+        )
+
+    @abstractmethod
+    def _standardize_single_inputs(
+        self,
+        data_folder: Path,
+        subject_id: str,
+        inputs: dict[str, Path | str],
+        subject_modality_separator: str,
+    ) -> None:
+        """
+        Standardize the input data to match the requirements of the selected algorithm.
+        """
+        pass
+
+    @abstractmethod
+    def _standardize_batch_inputs(
+        self, data_folder: Path, subjects: list[Path], input_name_schema: str
+    ) -> Dict[str, str]:
+        """
+        Standardize the input data to match the requirements of the selected algorithm.
+        """
+        pass
+
+    def extract_identifier_from_subject_id(self, subject_id: str) -> str:
+        """
+        Extract the index from the subject ID
+        Args:
+            subject_id (str): Subject ID of the input
+        Returns:
+            str: Extracted identifier
+        """
+        return "-".join(subject_id.split("-")[-2:])
+
+    def _process_single_output(
+        self,
+        tmp_output_folder: Path | str,
+        subject_id: str,
+        output_file: Path,
+    ) -> None:
+        """
+        Process the output of a single inference run and save it in the specified file.
+
+        Args:
+            tmp_output_folder (Path | str): Folder with the algorithm output
+            subject_id (str): Subject ID of the output
+            output_file (Path): Path to the desired output file
+        """
+        # rename 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__()
+        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:
+                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()
+        output_file.parent.mkdir(parents=True, exist_ok=True)
+        shutil.move(algorithm_output, output_file)
+
+    def _process_batch_output(
+        self,
+        tmp_output_folder: Path | str,
+        output_folder: Path,
+        mapping: dict[str, str],
+    ) -> None:
+        """
+        Process the outputs of a batch inference run and save them in the specified folder.
+
+        Args:
+            tmp_output_folder (Path | str): Folder with the algorithm outputs
+            output_folder (Path): Folder to save the outputs
+            mapping (dict[str, str]): Mapping from internal to external subject names
+        """
+        # move outputs and change name back to initially provided one
+        output_folder = Path(output_folder)
+        output_folder.mkdir(parents=True, exist_ok=True)
+        for internal_name, external_name in mapping.items():
+            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__()
+                )
+                try:
+                    modality = algorithm_output.name.split("-")[-1].split(".")[0]
+                except IndexError:
+                    logger.warning(
+                        f"Could not extract modality from {algorithm_output.name}"
+                    )
+                    modality = None
+                output_file = (
+                    output_folder
+                    / f"{external_name}{'-' + modality if modality else ''}.nii.gz"
+                )
+            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:
+                    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)
+
+    def _get_backend_runner(self, backend: Backends) -> Optional[Callable]:
+        backend_dispatch = {
+            Backends.DOCKER: run_docker_container,
+            Backends.SINGULARITY: run_singularity_container,
+        }
+        runner = backend_dispatch.get(backend, None)
+        return runner
+
+    def _infer_single(
+        self,
+        inputs: dict[str, Path | str],
+        output_file: Path | str,
+        log_file: Optional[Path | str] = None,
+        backend: Backends = Backends.DOCKER,
+    ) -> None:
+        """
+        Perform a single inference run with the provided inputs and save the output in the specified file.
+
+        Args:
+            inputs (dict[str, Path  |  str]): Input Images for the task
+            output_file (Path | str): File to save the output
+            log_file (Optional[Path  |  str], optional): Log file with extra information. Defaults to None.
+            backend (Backends | str, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+        """
+        with InferenceSetup(log_file=log_file) as (tmp_data_folder, tmp_output_folder):
+            logger.info(f"Performing single inference")
+
+            # the id here is arbitrary
+            subject_id = self.algorithm.run_args.input_name_schema.format(id=0)
+
+            self._standardize_single_inputs(
+                data_folder=tmp_data_folder,
+                subject_id=subject_id,
+                inputs=inputs,
+                subject_modality_separator=self.algorithm.run_args.subject_modality_separator,
+            )
+
+            runner = self._get_backend_runner(backend)
+            if runner is None:
+                raise ValueError(f"Unsupported backend: {backend}")
+            runner(
+                algorithm=self.algorithm,
+                data_path=tmp_data_folder,
+                output_path=tmp_output_folder,
+                cuda_devices=self.cuda_devices,
+                force_cpu=self.force_cpu,
+            )
+            self._process_single_output(
+                tmp_output_folder=tmp_output_folder,
+                subject_id=subject_id,
+                output_file=Path(output_file),
+            )
+            logger.info(f"Saved output to: {Path(output_file).absolute()}")
+
+    def _infer_batch(
+        self,
+        data_folder: Path | str,
+        output_folder: Path | str,
+        log_file: Optional[Path | str] = None,
+        backend: Backends = Backends.DOCKER,
+    ):
+        """Perform a batch inference run with the provided inputs and save the outputs in the specified folder.
+
+        Args:
+            data_folder (Path | str): Folder with the input data
+            output_folder (Path | str): Folder to save the outputs
+            log_file (Optional[Path  |  str], optional): Log file with extra information. Defaults to None.
+            backend (Backends, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+        """
+        with InferenceSetup(log_file=log_file) as (tmp_data_folder, tmp_output_folder):
+
+            # find subjects
+            subjects = [f for f in Path(data_folder).iterdir() if f.is_dir()]
+            logger.info(
+                f"Found {len(subjects)} subjects: {', '.join([s.name for s in subjects][:5])} {' ...' if len(subjects) > 5 else '' }"
+            )
+            # map to brats names
+            internal_external_name_map = self._standardize_batch_inputs(
+                data_folder=tmp_data_folder,
+                subjects=subjects,
+                input_name_schema=self.algorithm.run_args.input_name_schema,
+            )
+            logger.info(f"Standardized input names to match algorithm requirements.")
+            runner = self._get_backend_runner(backend)
+            if runner is None:
+                raise ValueError(f"Unsupported backend: {backend}")
+            # run inference in container
+            runner(
+                algorithm=self.algorithm,
+                data_path=tmp_data_folder,
+                output_path=tmp_output_folder,
+                cuda_devices=self.cuda_devices,
+                force_cpu=self.force_cpu,
+                internal_external_name_map=internal_external_name_map,
+            )
+
+            self._process_batch_output(
+                tmp_output_folder=tmp_output_folder,
+                output_folder=output_folder,
+                mapping=internal_external_name_map,
+            )
+
+            logger.info(f"Saved outputs to: {Path(output_folder).absolute()}")