Merge pull request #52 from BrainLesion/45-add-2024-segmentation-algo…

…rithms 45 add 2024 segmentation algorithms
BrainLesion · Dec 20, 2024 · e80390b · e80390b
2 parents e1569c6 + a1abdaa
commit e80390b
Show file tree

Hide file tree

Showing 22 changed files with 879 additions and 175 deletions.
diff --git a/README.md b/README.md
@@ -51,7 +51,7 @@ from brats import AdultGliomaSegmenter
 from brats.constants import AdultGliomaAlgorithms
 
 segmenter = AdultGliomaSegmenter(algorithm=AdultGliomaAlgorithms.BraTS23_1, cuda_devices="0")
-# these parameters are optional, by default the winning algorithm will be used on cuda:0
+# these parameters are optional, by default the winning algorithm of 2023 will be used on cuda:0
 segmenter.infer_single(
     t1c="path/to/t1c.nii.gz",
     t1n="path/to/t1n.nii.gz",
@@ -65,6 +65,9 @@ segmenter.infer_single(
 
 | Year | Rank | Author                            | Paper                                      | CPU Support | Key Enum                                                                                                             |
 | ---- | ---- | --------------------------------- | ------------------------------------------ | ----------- | -------------------------------------------------------------------------------------------------------------------- |
+| 2024 | 1st  | _André Ferreira, et al._          | N/A                                        | &#x274C;    | [BraTS24_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS24_1) |
+| 2024 | 2nd  | _Team kimbab_                     | N/A                                        | &#x274C;    | [BraTS24_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS24_2) |
+| 2024 | 3rd  | _Adrian Celaya_                   | N/A                                        | &#x2705;    | [BraTS24_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS24_3) |
 | 2023 | 1st  | _André Ferreira, et al._          | [Link](https://arxiv.org/abs/2402.17317v1) | &#x274C;    | [BraTS23_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS23_1) |
 | 2023 | 2nd  | _Andriy Myronenko, et al._        | N/A                                        | &#x274C;    | [BraTS23_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS23_2) |
 | 2023 | 3rd  | _Fadillah Adamsyah Maani, et al._ | N/A                                        | &#x274C;    | [BraTS23_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AdultGliomaAlgorithms.BraTS23_3) |
@@ -94,6 +97,9 @@ segmenter.infer_single(
 
 | Year | Rank | Author                     | Paper | CPU Support | Key Enum                                                                                                        |
 | ---- | ---- | -------------------------- | ----- | ----------- | --------------------------------------------------------------------------------------------------------------- |
+| 2024 | 1st  | _Zhifan Jiang et al._      | N/A   | &#x274C;    | [BraTS24_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS24_1) |
+| 2024 | 2nd  | _Long Bai, et al._         | N/A   | &#x2705;    | [BraTS24_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS24_2) |
+| 2024 | 1st  | _Sarim Hashmi, et al._     | N/A   | &#x274C;    | [BraTS24_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS24_3) |
 | 2023 | 1st  | _Andriy Myronenko, et al._ | TODO  | &#x274C;    | [BraTS23_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS23_1) |
 | 2023 | 2nd  | _Alyssa R Amod, et al._    | N/A   | &#x274C;    | [BraTS23_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS23_2) |
 | 2023 | 3rd  | _Ziyan Huang, et al._      | N/A   | &#x2705;    | [BraTS23_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.AfricaAlgorithms.BraTS23_3) |
@@ -104,25 +110,43 @@ segmenter.infer_single(
 <summary> <strong> Meningioma Segmentation </strong> </summary>
 <br>
 
+**Note**  
+Unlike other segmentation challenges the expected inputs for the Meningioma Segmentation Algorithms differ between years. 
+- _2023_: All 4 modalities are used (t1c, t1n, t2f, t2w)
+- _2024_: Only t1c is used  
+
+Therefore the usage differs slightly, depending on which algorithm is used. To understand why, please refer to the [2024 challenge manuscript](https://arxiv.org/abs/2405.18383).
+
 ```python
 from brats import MeningiomaSegmenter
 from brats.constants import MeningiomaAlgorithms
 
+### Example for 2023 algorithms
 segmenter = MeningiomaSegmenter(algorithm=MeningiomaAlgorithms.BraTS23_1, cuda_devices="0")
 # these parameters are optional, by default the winning algorithm will be used on cuda:0
 segmenter.infer_single(
     t1c="path/to/t1c.nii.gz",
     t1n="path/to/t1n.nii.gz",
     t2f="path/to/t2f.nii.gz",
     t2w="path/to/t2w.nii.gz",
-    output_file="segmentation.nii.gz",
+    output_file="segmentation_23.nii.gz",
+)
+
+### Example for 2024 algorithms
+segmenter = MeningiomaSegmenter(algorithm=MeningiomaAlgorithms.BraTS24_1, cuda_devices="0")
+segmenter.infer_single(
+    t1c="path/to/t1c.nii.gz",
+    output_file="segmentation_24.nii.gz",
 )
 ```
 
 **Class:** `brats.MeningiomaSegmenter` ([Docs](https://brats.readthedocs.io/en/latest/core/segmentation_algorithms.html#brats.core.segmentation_algorithms.MeningiomaSegmenter))
 
 | Year | Rank | Author                     | Paper | CPU Support | Key Enum                                                                                                            |
 | ---- | ---- | -------------------------- | ----- | ----------- | ------------------------------------------------------------------------------------------------------------------- |
+| 2024 | 1st  | _Valeria Abramova_         | N/A   | &#x274C;    | [BraTS24_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS24_1) |
+| 2024 | 2nd  | _Mehdi Astaraki_           | N/A   | &#x274C;    | [BraTS24_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS24_2) |
+| 2024 | 3rd  | _Andre Ferreira, et al._   | N/A   | &#x2705;    | [BraTS24_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS24_3) |
 | 2023 | 1st  | _Andriy Myronenko, et al._ | N/A   | &#x274C;    | [BraTS23_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS23_1) |
 | 2023 | 2nd  | _Ziyan Huang, et al._      | N/A   | &#x2705;    | [BraTS23_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS23_2) |
 | 2023 | 3rd  | _Zhifan Jiang et al._      | N/A   | &#x274C;    | [BraTS23_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.MeningiomaAlgorithms.BraTS23_3) |
@@ -181,6 +205,9 @@ segmenter.infer_single(
 
 | Year | Rank | Author                     | Paper | CPU Support | Key Enum                                                                                                           |
 | ---- | ---- | -------------------------- | ----- | ----------- | ------------------------------------------------------------------------------------------------------------------ |
+| 2024 | 1st  | _Tim Mulvany, et al._      | N/A   | &#x274C;    | [BraTS24_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS24_1) |
+| 2024 | 2nd  | _Mehdi Astaraki_           | N/A   | &#x274C;    | [BraTS24_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS24_2) |
+| 2024 | 3rd  | _Sarim Hashmi, et al._     | N/A   | &#x274C;    | [BraTS24_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS24_3) |
 | 2023 | 1st  | _Zhifan Jiang et al._      | N/A   | &#x274C;    | [BraTS23_1](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS23_1) |
 | 2023 | 2nd  | _Andriy Myronenko, et al._ | N/A   | &#x274C;    | [BraTS23_2](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS23_2) |
 | 2023 | 3rd  | _Yubo Zhou_                | N/A   | &#x274C;    | [BraTS23_3](https://brats.readthedocs.io/en/latest/utils/utils.html#brats.constants.PediatricAlgorithms.BraTS23_3) |

diff --git a/brats/constants.py b/brats/constants.py
@@ -29,6 +29,13 @@ class Algorithms(str, Enum):
 class AdultGliomaAlgorithms(Algorithms):
     """Constants for the available adult glioma segmentation algorithms."""
 
+    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 """
+
     BraTS23_1 = "BraTS23_1"
     """BraTS23 Adult Glioma Segmentation 1st place (GPU only)"""
     BraTS23_2 = "BraTS23_2"
@@ -40,6 +47,13 @@ class AdultGliomaAlgorithms(Algorithms):
 class MeningiomaAlgorithms(Algorithms):
     """Constants for the available meningioma segmentation algorithms."""
 
+    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 """
+
     BraTS23_1 = "BraTS23_1"
     """BraTS23 Meningioma Segmentation 1st place (GPU only)"""
     BraTS23_2 = "BraTS23_2"
@@ -51,6 +65,13 @@ class MeningiomaAlgorithms(Algorithms):
 class PediatricAlgorithms(Algorithms):
     """Constants for the available pediatric segmentation algorithms."""
 
+    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"
@@ -62,6 +83,13 @@ class PediatricAlgorithms(Algorithms):
 class AfricaAlgorithms(Algorithms):
     """Constants for the available africa segmentation algorithms."""
 
+    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"
@@ -90,6 +118,7 @@ class InpaintingAlgorithms(Algorithms):
     """ BraTS24  Inpainting 2nd place """
     BraTS24_3 = "BraTS24_3"
     """ BraTS24  Inpainting 3rd place """
+
     BraTS23_1 = "BraTS23_1"
     """ BraTS23  Inpainting 1st place """
     BraTS23_2 = "BraTS23_2"

diff --git a/brats/core/brats_algorithm.py b/brats/core/brats_algorithm.py
@@ -4,7 +4,7 @@
 import sys
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Optional
+from typing import Dict, Optional
 
 from loguru import logger
 
@@ -51,7 +51,11 @@ def __init__(
 
     @abstractmethod
     def _standardize_single_inputs(
-        self, data_folder: Path, subject_id: str, inputs: dict[str, Path | str]
+        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.
@@ -61,7 +65,7 @@ def _standardize_single_inputs(
     @abstractmethod
     def _standardize_batch_inputs(
         self, data_folder: Path, subjects: list[Path], input_name_schema: str
-    ) -> None:
+    ) -> Dict[str, str]:
         """
         Standardize the input data to match the requirements of the selected algorithm.
         """
@@ -157,6 +161,7 @@ def _infer_single(
                 data_folder=tmp_data_folder,
                 subject_id=subject_id,
                 inputs=inputs,
+                subject_modality_separator=self.algorithm.run_args.subject_modality_separator,
             )
 
             run_container(
@@ -208,6 +213,7 @@ def _infer_batch(
                 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(

diff --git a/brats/core/docker.py b/brats/core/docker.py
@@ -124,11 +124,13 @@ def _get_additional_files_path(algorithm: AlgorithmData) -> Path:
     Returns:
         Path to the additional files
     """
-    # ensure weights are present and get path
-    if algorithm.weights is not None:
-        return check_additional_files_path(record_id=algorithm.weights.record_id)
+    # ensure additional_files are present and get path
+    if algorithm.additional_files is not None:
+        return check_additional_files_path(
+            record_id=algorithm.additional_files.record_id
+        )
     else:
-        # if no weights are directly specified a dummy weights folder will be mounted
+        # if no additional_files are directly specified a dummy additional_files folder will be mounted
         return get_dummy_path()
 
 
@@ -195,12 +197,12 @@ def _build_args(
     """
     # Build command that will be run in the docker container
     command_args = f"--data_path=/mlcube_io0 --output_path=/mlcube_io2"
-    if algorithm.weights is not None:
-        for i, param in enumerate(algorithm.weights.param_name):
-            weights_arg = f"--{param}=/mlcube_io1"
-            if algorithm.weights.checkpoint_path:
-                weights_arg += f"/{algorithm.weights.checkpoint_path[i]}"
-            command_args += f" {weights_arg}"
+    if algorithm.additional_files is not None:
+        for i, param in enumerate(algorithm.additional_files.param_name):
+            additional_files_arg = f"--{param}=/mlcube_io1"
+            if algorithm.additional_files.param_path:
+                additional_files_arg += f"/{algorithm.additional_files.param_path[i]}"
+            command_args += f" {additional_files_arg}"
 
     # Add parameters file arg if required
     params_arg = _get_parameters_arg(algorithm=algorithm)
@@ -245,14 +247,18 @@ def _observe_docker_output(container: docker.models.containers.Container) -> str
 
 
 def _sanity_check_output(
-    data_path: Path, output_path: Path, container_output: str
+    data_path: Path,
+    output_path: Path,
+    container_output: str,
+    internal_external_name_map: Optional[Dict[str, str]] = None,
 ) -> None:
     """Sanity check that the number of output files matches the number of input files and the output is not empty.
 
     Args:
         data_path (Path): The path to the input data
         output_path (Path): The path to the output data
         container_output (str): The output of the docker container
+        internal_external_name_map (Optional[Dict[str, str]]): Dictionary mapping internal name (in standardized format) to external subject name provided by user (only used for batch inference)
 
     Raises:
         BraTSContainerException: If not enough output files exist
@@ -262,7 +268,6 @@ def _sanity_check_output(
     # (should result in only counting actual inputs)
     inputs = [e for e in data_path.iterdir() if e.name.startswith("BraTS")]
     outputs = list(output_path.iterdir())
-
     if len(outputs) < len(inputs):
         logger.error(f"Docker container output: \n\r{container_output}")
         raise BraTSContainerException(
@@ -272,8 +277,18 @@ def _sanity_check_output(
     for i, output in enumerate(outputs, start=1):
         content = nib.load(output).get_fdata()
         if np.count_nonzero(content) == 0:
+            name = ""
+            if internal_external_name_map is not None:
+                name_key = [
+                    k
+                    for k in internal_external_name_map.keys()
+                    if output.name.startswith(k)
+                ]
+                if name_key:
+                    name = internal_external_name_map[name_key[0]]
+
             logger.warning(
-                f"""Output file {i} contains only zeros.
+                f"""Output file for subject {name + " "}contains only zeros.
                 Potentially the selected algorithm might not work properly with your data unless this behavior is correct for your use case.
                 If this seems wrong please try to use one of the other provided algorithms and file an issue on GitHub if the problem persists."""
             )
@@ -286,7 +301,7 @@ def _log_algorithm_info(algorithm: AlgorithmData):
         algorithm (AlgorithmData): algorithm data
     """
     logger.opt(colors=True).info(
-        f"Running algorithm: <light-green>{algorithm.meta.challenge} [{algorithm.meta.rank} place]</>"
+        f"Running algorithm: <light-green> BraTS {algorithm.meta.year} {algorithm.meta.challenge} [{algorithm.meta.rank} place]</>"
     )
     logger.opt(colors=True).info(
         f"<blue>(Paper)</> Consider citing the corresponding paper: {algorithm.meta.paper} by {algorithm.meta.authors}"
@@ -300,6 +315,7 @@ def run_container(
     output_path: Path,
     cuda_devices: str,
     force_cpu: bool,
+    internal_external_name_map: Optional[Dict[str, str]] = None,
 ):
     """Run a docker container for the provided algorithm.
 
@@ -309,6 +325,7 @@ def run_container(
         output_path (Path | str): The path to save the output
         cuda_devices (str): The CUDA devices to use
         force_cpu (bool): Whether to force CPU execution
+        internal_external_name_map (Dict[str, str]): Dictionary mapping internal name (in standardized format) to external subject name provided by user (only used for batch inference)
     """
     _log_algorithm_info(algorithm=algorithm)
 
@@ -353,7 +370,10 @@ def run_container(
     )
     container_output = _observe_docker_output(container=container)
     _sanity_check_output(
-        data_path=data_path, output_path=output_path, container_output=container_output
+        data_path=data_path,
+        output_path=output_path,
+        container_output=container_output,
+        internal_external_name_map=internal_external_name_map,
     )
 
     logger.debug(f"Docker container output: \n\r{container_output}")