BraTS 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,51 @@
1
+ Metadata-Version: 2.1
2
+ Name: BraTS
3
+ Version: 0.0.1
4
+ Summary: A small example package
5
+ Author-email: Example Author <author@example.com>
6
+ Project-URL: Homepage, https://github.com/pypa/sampleproject
7
+ Project-URL: Issues, https://github.com/pypa/sampleproject/issues
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+
14
+ [![PyPI version brats-algorithms](https://badge.fury.io/py/brats-algorithms.svg)](https://pypi.python.org/pypi/brats-algorithms/)
15
+ [![Documentation Status](https://readthedocs.org/projects/brats-algorithms/badge/?version=latest)](http://brats-algorithms.readthedocs.io/?badge=latest)
16
+ [![tests](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml/badge.svg)](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml)
17
+
18
+ # brats-algorithms
19
+
20
+ Top performing algorithms from the Brain Tumor Segmentation (BraTS) challenges, providing state-of-the-art models for brain lesion segmentation.
21
+
22
+ ## Features
23
+
24
+ - Access to top-performing algorithms from recent BraTS challenges.
25
+ - Easy-to-use inference API.
26
+ - Extensive documentation and examples. (TODO!)
27
+
28
+ ## Installation
29
+
30
+ With a Python 3.10+ environment, you can install `PACKAGE_NAME` directly from [PyPI](https://pypi.org/project/PACKAGE_NAME/):
31
+
32
+ ```sh
33
+ pip install PACKAGE_NAME
34
+ ```
35
+
36
+ ### Docker and NVIDIA Container Toolkit Setup
37
+
38
+ To run `PACKAGE_NAME` in a Docker container with GPU support (required for most algorithms), ensure you have Docker and NVIDIA Docker installed. Instructions:
39
+ - **Docker**: Installation instructions on the official [website](https://docs.docker.com/get-docker/)
40
+ - **NVIDIA Container Toolkit**: Installation instructions on the official [GitHub page](https://github.com/NVIDIA/nvidia-container-toolkit)
41
+
42
+ ## Use Cases and Tutorials
43
+
44
+ *TODO*
45
+
46
+ ## Contributing
47
+
48
+ We welcome contributions from the community, including bug reports, feature requests, and code contributions. For more information on how to contribute, please see our [CONTRIBUTING.md](CONTRIBUTING.md) file.
49
+
50
+ ## License
51
+ *TODO*
@@ -0,0 +1,16 @@
1
+ README.md
2
+ pyproject.toml
3
+ BraTS.egg-info/PKG-INFO
4
+ BraTS.egg-info/SOURCES.txt
5
+ BraTS.egg-info/dependency_links.txt
6
+ BraTS.egg-info/top_level.txt
7
+ brats/__init__.py
8
+ brats/algorithm_config.py
9
+ brats/algorithms.py
10
+ brats/constants.py
11
+ brats/docker.py
12
+ brats/exceptions.py
13
+ brats/utils.py
14
+ brats/weights.py
15
+ tests/test_algorithm_config.py
16
+ tests/test_utils.py
@@ -0,0 +1 @@
1
+ brats
brats-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,51 @@
1
+ Metadata-Version: 2.1
2
+ Name: BraTS
3
+ Version: 0.0.1
4
+ Summary: A small example package
5
+ Author-email: Example Author <author@example.com>
6
+ Project-URL: Homepage, https://github.com/pypa/sampleproject
7
+ Project-URL: Issues, https://github.com/pypa/sampleproject/issues
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+
14
+ [![PyPI version brats-algorithms](https://badge.fury.io/py/brats-algorithms.svg)](https://pypi.python.org/pypi/brats-algorithms/)
15
+ [![Documentation Status](https://readthedocs.org/projects/brats-algorithms/badge/?version=latest)](http://brats-algorithms.readthedocs.io/?badge=latest)
16
+ [![tests](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml/badge.svg)](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml)
17
+
18
+ # brats-algorithms
19
+
20
+ Top performing algorithms from the Brain Tumor Segmentation (BraTS) challenges, providing state-of-the-art models for brain lesion segmentation.
21
+
22
+ ## Features
23
+
24
+ - Access to top-performing algorithms from recent BraTS challenges.
25
+ - Easy-to-use inference API.
26
+ - Extensive documentation and examples. (TODO!)
27
+
28
+ ## Installation
29
+
30
+ With a Python 3.10+ environment, you can install `PACKAGE_NAME` directly from [PyPI](https://pypi.org/project/PACKAGE_NAME/):
31
+
32
+ ```sh
33
+ pip install PACKAGE_NAME
34
+ ```
35
+
36
+ ### Docker and NVIDIA Container Toolkit Setup
37
+
38
+ To run `PACKAGE_NAME` in a Docker container with GPU support (required for most algorithms), ensure you have Docker and NVIDIA Docker installed. Instructions:
39
+ - **Docker**: Installation instructions on the official [website](https://docs.docker.com/get-docker/)
40
+ - **NVIDIA Container Toolkit**: Installation instructions on the official [GitHub page](https://github.com/NVIDIA/nvidia-container-toolkit)
41
+
42
+ ## Use Cases and Tutorials
43
+
44
+ *TODO*
45
+
46
+ ## Contributing
47
+
48
+ We welcome contributions from the community, including bug reports, feature requests, and code contributions. For more information on how to contribute, please see our [CONTRIBUTING.md](CONTRIBUTING.md) file.
49
+
50
+ ## License
51
+ *TODO*
brats-0.0.1/README.md ADDED
@@ -0,0 +1,38 @@
1
+ [![PyPI version brats-algorithms](https://badge.fury.io/py/brats-algorithms.svg)](https://pypi.python.org/pypi/brats-algorithms/)
2
+ [![Documentation Status](https://readthedocs.org/projects/brats-algorithms/badge/?version=latest)](http://brats-algorithms.readthedocs.io/?badge=latest)
3
+ [![tests](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml/badge.svg)](https://github.com/BrainLesion/brats-algorithms/actions/workflows/tests.yml)
4
+
5
+ # brats-algorithms
6
+
7
+ Top performing algorithms from the Brain Tumor Segmentation (BraTS) challenges, providing state-of-the-art models for brain lesion segmentation.
8
+
9
+ ## Features
10
+
11
+ - Access to top-performing algorithms from recent BraTS challenges.
12
+ - Easy-to-use inference API.
13
+ - Extensive documentation and examples. (TODO!)
14
+
15
+ ## Installation
16
+
17
+ With a Python 3.10+ environment, you can install `PACKAGE_NAME` directly from [PyPI](https://pypi.org/project/PACKAGE_NAME/):
18
+
19
+ ```sh
20
+ pip install PACKAGE_NAME
21
+ ```
22
+
23
+ ### Docker and NVIDIA Container Toolkit Setup
24
+
25
+ To run `PACKAGE_NAME` in a Docker container with GPU support (required for most algorithms), ensure you have Docker and NVIDIA Docker installed. Instructions:
26
+ - **Docker**: Installation instructions on the official [website](https://docs.docker.com/get-docker/)
27
+ - **NVIDIA Container Toolkit**: Installation instructions on the official [GitHub page](https://github.com/NVIDIA/nvidia-container-toolkit)
28
+
29
+ ## Use Cases and Tutorials
30
+
31
+ *TODO*
32
+
33
+ ## Contributing
34
+
35
+ We welcome contributions from the community, including bug reports, feature requests, and code contributions. For more information on how to contribute, please see our [CONTRIBUTING.md](CONTRIBUTING.md) file.
36
+
37
+ ## License
38
+ *TODO*
@@ -0,0 +1,5 @@
1
+ from brats.algorithms import (
2
+ AdultGliomaSegmenter,
3
+ MeningiomaSegmenter,
4
+ PediatricSegmenter,
5
+ )
@@ -0,0 +1,90 @@
1
+ from dataclasses import dataclass
2
+ from pathlib import Path
3
+ from typing import Dict, Optional
4
+ from brats.exceptions import AlgorithmConfigException
5
+
6
+ import yaml
7
+ from dacite import DaciteError, from_dict
8
+
9
+
10
+ @dataclass
11
+ class MetaData:
12
+ """Dataclass for the meta data"""
13
+
14
+ authors: str
15
+ """The authors of the algorithm"""
16
+ paper: str
17
+ """If available, a url to the paper of the algorithm"""
18
+ challenge: str
19
+ """The challenge the algorithm was submitted to"""
20
+
21
+
22
+ @dataclass
23
+ class RunArgs:
24
+ """Dataclass for the run arguments"""
25
+
26
+ docker_image: str
27
+ """The Docker image containing the algorithm"""
28
+ input_name_schema: str
29
+ """The input name schema for the algorithm"""
30
+ parameters_file: bool
31
+ """Whether the algorithm requires a parameters file"""
32
+ requires_root: bool
33
+ """Whether the Docker container requires root access. This is !discouraged! but some submission do not work without it"""
34
+ shm_size: Optional[str] = "2gb"
35
+ """The required shared memory size for the Docker container"""
36
+ cpu_compatible: Optional[bool] = False
37
+ """Whether the algorithm is compatible with CPU"""
38
+
39
+
40
+ @dataclass
41
+ class WeightsData:
42
+ """Dataclass for the weights data"""
43
+
44
+ record_id: str
45
+ """The Zenodo record ID of the weights"""
46
+ param_name: Optional[str] = "weights"
47
+ """The parameter that specifies the weights folder in the algorithm execution, typically 'weights' but differs for some"""
48
+
49
+
50
+ @dataclass
51
+ class AlgorithmData:
52
+ """Dataclass for the algorithm data"""
53
+
54
+ meta: MetaData
55
+ """The meta data of the algorithm"""
56
+ run_args: RunArgs
57
+ """The run arguments of the algorithm"""
58
+ weights: Optional[WeightsData]
59
+ """The weights data of the algorithm. Optional since some algorithms include weights in the docker image"""
60
+
61
+
62
+ @dataclass
63
+ class AlgorithmList:
64
+ algorithms: Dict[str, AlgorithmData]
65
+
66
+
67
+ def load_algorithms(file_path: Path) -> Dict[str, AlgorithmData]:
68
+ """Load the algorithms data from the specified yaml file
69
+
70
+ Params:
71
+ file_path (str): The path to the yaml file
72
+
73
+ Raises:
74
+ FileNotFoundError: If the file is not found
75
+
76
+ Returns:
77
+ Dict[str, AlgorithmData]: Dict of algorithm @AlgorithmKeys:@AlgorithmData pairs
78
+ """
79
+ try:
80
+ with open(file_path, "r") as file:
81
+ data = yaml.safe_load(file)
82
+ except FileNotFoundError:
83
+ raise FileNotFoundError("Algorithm meta data file not found")
84
+
85
+ try:
86
+ # Convert the dictionary to the dataclass
87
+ algorithms = from_dict(data_class=AlgorithmList, data=data).algorithms
88
+ except DaciteError as e:
89
+ raise AlgorithmConfigException(f"Error loading algorithm data: {e}")
90
+ return algorithms
@@ -0,0 +1,262 @@
1
+ from __future__ import annotations
2
+
3
+ import shutil
4
+ import tempfile
5
+ from pathlib import Path
6
+ from typing import Optional
7
+
8
+ from loguru import logger
9
+
10
+ from brats.algorithm_config import load_algorithms
11
+ from brats.constants import (
12
+ ADULT_GLIOMA_SEGMENTATION_ALGORITHMS,
13
+ MENINGIOMA_SEGMENTATION_ALGORITHMS,
14
+ PEDIATRIC_SEGMENTATION_ALGORITHMS,
15
+ AdultGliomaAlgorithms,
16
+ Algorithms,
17
+ MeningiomaAlgorithms,
18
+ PediatricAlgorithms,
19
+ )
20
+ from brats.docker import run_docker
21
+ from brats.utils import (
22
+ standardize_subject_inputs,
23
+ standardize_subjects_inputs_list,
24
+ )
25
+
26
+
27
+ class BraTSAlgorithm:
28
+ def __init__(
29
+ self,
30
+ algorithm: Algorithms,
31
+ algorithms_file_path: Path,
32
+ cuda_devices: Optional[str] = "0",
33
+ force_cpu: bool = False,
34
+ ):
35
+ # inference device setup
36
+ self.force_cpu = force_cpu
37
+ self.cuda_devices = cuda_devices
38
+
39
+ self.algorithm_list = load_algorithms(file_path=algorithms_file_path)
40
+ # save algorithm identifier for logging etc.
41
+ self.algorithm_key = algorithm.value
42
+ # data for selected algorithm
43
+ self.algorithm = self.algorithm_list[algorithm.value]
44
+
45
+ logger.info(
46
+ f"Instantiated {self.__class__.__name__} with algorithm: {self.algorithm_key} by {self.algorithm.meta.authors}"
47
+ )
48
+
49
+ def _log_algorithm_info(self):
50
+ """Log information about the selected algorithm."""
51
+ logger.opt(colors=True).info(
52
+ f"Running algorithm: <light-green>{self.algorithm.meta.challenge}</>"
53
+ )
54
+ logger.opt(colors=True).info(
55
+ f"<blue>(Docker image)</>: {self.algorithm.run_args.docker_image}"
56
+ )
57
+ logger.opt(colors=True).info(
58
+ f"<blue>(Paper)</> Consider citing the corresponding paper: {self.algorithm.meta.paper} by {self.algorithm.meta.authors}"
59
+ )
60
+
61
+ def infer_single(
62
+ self,
63
+ t1c: Path | str,
64
+ t1n: Path | str,
65
+ t2f: Path | str,
66
+ t2w: Path | str,
67
+ output_file: Path | str,
68
+ log_file: Optional[Path | str] = None,
69
+ ):
70
+ """Perform inference on a single subject with the provided images and save the result to the output file.
71
+
72
+ Args:
73
+ t1c (Path | str): Path to the T1c image
74
+ t1n (Path | str): Path to the T1n image
75
+ t2f (Path | str): Path to the T2f image
76
+ t2w (Path | str): Path to the T2w image
77
+ output_file (Path | str): Path to save the segmentation
78
+ log_file (Path | str, optional): Save logs to this file
79
+ """
80
+ # setup temp input folder with the provided images
81
+ temp_data_folder = Path(tempfile.mkdtemp())
82
+ temp_output_folder = Path(tempfile.mkdtemp())
83
+ if log_file is not None:
84
+ inference_log_file = logger.add(log_file, level="INFO")
85
+ logger.info(f"Logging to: {log_file}")
86
+
87
+ try:
88
+ logger.info(f"Performing single inference ")
89
+
90
+ # the id here is arbitrary
91
+ subject_id = self.algorithm.run_args.input_name_schema.format(id=0)
92
+
93
+ standardize_subject_inputs(
94
+ data_folder=temp_data_folder,
95
+ subject_id=subject_id,
96
+ t1c=t1c,
97
+ t1n=t1n,
98
+ t2f=t2f,
99
+ t2w=t2w,
100
+ )
101
+
102
+ self._log_algorithm_info()
103
+ run_docker(
104
+ algorithm=self.algorithm,
105
+ data_path=temp_data_folder,
106
+ output_path=temp_output_folder,
107
+ cuda_devices=self.cuda_devices,
108
+ force_cpu=self.force_cpu,
109
+ )
110
+ # rename output
111
+ segmentation = Path(temp_output_folder) / f"{subject_id}.nii.gz"
112
+
113
+ # ensure path exists and rename output to the desired path
114
+ output_file = Path(output_file).absolute()
115
+ output_file.parent.mkdir(parents=True, exist_ok=True)
116
+ shutil.move(segmentation, output_file)
117
+ logger.info(f"Saved segmentation to: {output_file.absolute()}")
118
+
119
+ finally:
120
+ shutil.rmtree(temp_data_folder)
121
+ shutil.rmtree(temp_output_folder)
122
+ if log_file is not None:
123
+ logger.remove(inference_log_file)
124
+
125
+ def infer_batch(
126
+ self,
127
+ data_folder: Path | str,
128
+ output_folder: Path | str,
129
+ log_file: Optional[Path | str] = None,
130
+ ):
131
+ """Perform segmentation on a batch of subjects with the provided images and save the results to the output folder. \n
132
+ Requires the following structure:\n
133
+ data_folder\n
134
+ ┣ A\n
135
+ ┃ ┣ A-t1c.nii.gz\n
136
+ ┃ ┣ A-t1n.nii.gz\n
137
+ ┃ ┣ A-t2f.nii.gz\n
138
+ ┃ ┗ A-t2w.nii.gz\n
139
+ ┣ B\n
140
+ ┃ ┣ B-t1c.nii.gz\n
141
+ ┃ ┣ ...\n
142
+
143
+
144
+ Args:
145
+ data_folder (Path | str): Folder containing the subjects with required structure
146
+ output_folder (Path | str): Output folder to save the segmentations
147
+ log_file (Path | str, optional): Save logs to this file
148
+ """
149
+
150
+ data_folder = Path(data_folder)
151
+ output_folder = Path(output_folder)
152
+
153
+ temp_data_folder = Path(tempfile.mkdtemp())
154
+ temp_output_folder = Path(tempfile.mkdtemp())
155
+ if log_file:
156
+ inference_log_file = logger.add(log_file, level="INFO", catch=True)
157
+ logger.info(f"Logging to: {log_file.absolute()}")
158
+ try:
159
+ self._log_algorithm_info()
160
+ # find subjects
161
+ subjects = [f for f in data_folder.iterdir() if f.is_dir()]
162
+ logger.info(
163
+ f"Found {len(subjects)} subjects: {', '.join([s.name for s in subjects][:5])} {' ...' if len(subjects) > 5 else '' }"
164
+ )
165
+ # map to brats names
166
+ internal_external_name_map = standardize_subjects_inputs_list(
167
+ subjects=subjects,
168
+ temp_data_folder=temp_data_folder,
169
+ input_name_schema=self.algorithm.run_args.input_name_schema,
170
+ )
171
+ logger.info(f"Standardized input names to match algorithm requirements.")
172
+
173
+ # run inference in container
174
+ run_docker(
175
+ algorithm=self.algorithm,
176
+ data_path=temp_data_folder,
177
+ output_path=temp_output_folder,
178
+ cuda_devices=self.cuda_devices,
179
+ force_cpu=self.force_cpu,
180
+ )
181
+
182
+ # move outputs and change name back to initially provided one
183
+ for internal_name, external_name in internal_external_name_map.items():
184
+ segmentation = Path(temp_output_folder) / f"{internal_name}.nii.gz"
185
+ output_file = output_folder / f"{external_name}.nii.gz"
186
+ shutil.move(segmentation, output_file)
187
+
188
+ logger.info(f"Saved results to: {output_folder.absolute()}")
189
+ finally:
190
+ shutil.rmtree(temp_data_folder)
191
+ shutil.rmtree(temp_output_folder)
192
+ if log_file:
193
+ logger.remove(inference_log_file)
194
+
195
+
196
+ class AdultGliomaSegmenter(BraTSAlgorithm):
197
+
198
+ def __init__(
199
+ self,
200
+ algorithm: AdultGliomaAlgorithms = AdultGliomaAlgorithms.BraTS23_1,
201
+ cuda_devices: Optional[str] = "0",
202
+ force_cpu: bool = False,
203
+ ):
204
+ """Provides algorithms to perform tumor segmentation on adult glioma MRI data.
205
+
206
+ Args:
207
+ algorithm (AdultGliomaAlgorithms, optional): Select an algorithm. Defaults to AdultGliomaAlgorithms.BraTS23_1.
208
+ cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0".
209
+ force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False.
210
+ """
211
+ super().__init__(
212
+ algorithm=algorithm,
213
+ algorithms_file_path=ADULT_GLIOMA_SEGMENTATION_ALGORITHMS,
214
+ cuda_devices=cuda_devices,
215
+ force_cpu=force_cpu,
216
+ )
217
+
218
+
219
+ class MeningiomaSegmenter(BraTSAlgorithm):
220
+
221
+ def __init__(
222
+ self,
223
+ algorithm: MeningiomaAlgorithms = MeningiomaAlgorithms.BraTS23_1,
224
+ cuda_devices: Optional[str] = "0",
225
+ force_cpu: bool = False,
226
+ ):
227
+ """Provides algorithms to perform tumor segmentation on adult meningioma MRI data.
228
+
229
+ Args:
230
+ algorithm (AdultGliomaAlgorithms, optional): Select an algorithm. Defaults to MeningiomaAlgorithms.BraTS23_1.
231
+ cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0".
232
+ force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False.
233
+ """
234
+ super().__init__(
235
+ algorithm=algorithm,
236
+ algorithms_file_path=MENINGIOMA_SEGMENTATION_ALGORITHMS,
237
+ cuda_devices=cuda_devices,
238
+ force_cpu=force_cpu,
239
+ )
240
+
241
+
242
+ class PediatricSegmenter(BraTSAlgorithm):
243
+
244
+ def __init__(
245
+ self,
246
+ algorithm: PediatricAlgorithms = PediatricAlgorithms.BraTS23_1,
247
+ cuda_devices: Optional[str] = "0",
248
+ force_cpu: bool = False,
249
+ ):
250
+ """Provides algorithms to perform tumor segmentation on pediatric MRI data
251
+
252
+ Args:
253
+ algorithm (AdultGliomaAlgorithms, optional): Select an algorithm. Defaults to PediatricAlgorithms.BraTS23_1.
254
+ cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0".
255
+ force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False.
256
+ """
257
+ super().__init__(
258
+ algorithm=algorithm,
259
+ algorithms_file_path=PEDIATRIC_SEGMENTATION_ALGORITHMS,
260
+ cuda_devices=cuda_devices,
261
+ force_cpu=force_cpu,
262
+ )
@@ -0,0 +1,48 @@
1
+ from enum import Enum
2
+ from pathlib import Path
3
+
4
+
5
+ class Algorithms(str, Enum):
6
+ """Parent class for constants of the available algorithms."""
7
+
8
+ pass
9
+
10
+
11
+ class AdultGliomaAlgorithms(Algorithms):
12
+ """Constants for the available adult glioma segmentation algorithms."""
13
+
14
+ BraTS23_1 = "BraTS23_1"
15
+ """BraTS23 Adult Glioma Segmentation 1st place (GPU only)"""
16
+ BraTS23_2 = "BraTS23_2"
17
+ """BraTS23 Adult Glioma Segmentation 2nd place (GPU only)"""
18
+ BraTS23_3 = "BraTS23_3"
19
+ """BraTS23 Adult Glioma Segmentation 3rd place (GPU only)"""
20
+
21
+
22
+ class MeningiomaAlgorithms(Algorithms):
23
+ """Constants for the available meningioma segmentation algorithms."""
24
+
25
+ BraTS23_1 = "BraTS23_1"
26
+ """BraTS23 Meningioma Segmentation 1st place (GPU only)"""
27
+ BraTS23_2 = "BraTS23_2"
28
+ """BraTS23 Meningioma Segmentation 2nd place (GPU and CPU)"""
29
+ BraTS23_3 = "BraTS23_3"
30
+ """BraTS23 Meningioma Segmentation 3rd place (GPU only)"""
31
+
32
+
33
+ class PediatricAlgorithms(Algorithms):
34
+ """Constants for the available pediatric segmentation algorithms."""
35
+
36
+ BraTS23_1 = "BraTS23_1"
37
+ """BraTS23 Pediatric Segmentation 1st place (GPU only)"""
38
+ BraTS23_2 = "BraTS23_2"
39
+ """BraTS23 Pediatric Segmentation 2nd place (GPU only)"""
40
+ BraTS23_3 = "BraTS23_3"
41
+ """BraTS23 Pediatric Segmentation 3rd place (GPU only)"""
42
+
43
+
44
+ # meta data file paths
45
+ ALGORITHM_DIR = Path(__file__).parent / "algorithms"
46
+ ADULT_GLIOMA_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "adult_glioma.yml"
47
+ MENINGIOMA_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "meningioma.yml"
48
+ PEDIATRIC_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "pediatric.yml"
@@ -0,0 +1,261 @@
1
+ from __future__ import annotations
2
+
3
+ import os
4
+ import subprocess
5
+ from pathlib import Path
6
+ from typing import Dict, List, Tuple
7
+
8
+ import docker
9
+ from halo import Halo
10
+ from loguru import logger
11
+ from rich.progress import Progress
12
+
13
+ from brats.algorithm_config import AlgorithmData
14
+ from brats.exceptions import AlgorithmNotCPUCompatibleException
15
+ from brats.weights import check_model_weights, get_dummy_weights_path
16
+
17
+ client = docker.from_env()
18
+
19
+
20
+ def _show_docker_pull_progress(tasks: Dict, progress: Progress, line: Dict):
21
+ """Show the progress of a docker pull operation.
22
+
23
+ Args:
24
+ tasks (Dict): The tasks to update
25
+ progress (Progress): The progress bar to update
26
+ line (Dict): the next line from docker.client.api.pull stream
27
+ """
28
+ if line["status"] == "Downloading":
29
+ task_key = f'[Download {line["id"]}]'
30
+ elif line["status"] == "Extracting":
31
+ task_key = f'[Extract {line["id"]}]'
32
+ else:
33
+ return
34
+
35
+ if task_key not in tasks.keys():
36
+ tasks[task_key] = progress.add_task(
37
+ f"{task_key}", total=line["progressDetail"]["total"]
38
+ )
39
+ else:
40
+ progress.update(tasks[task_key], completed=line["progressDetail"]["current"])
41
+
42
+
43
+ def _ensure_image(image: str):
44
+ """Ensure the docker image is present on the system. If not, pull it.
45
+
46
+ Args:
47
+ image (str): The docker image to pull
48
+ """
49
+ if not client.images.list(name=image):
50
+ logger.info(f"Pulling docker image {image}")
51
+ tasks = {}
52
+ with Progress() as progress:
53
+ resp = client.api.pull(image, stream=True, decode=True)
54
+ for line in resp:
55
+ _show_docker_pull_progress(tasks=tasks, progress=progress, line=line)
56
+
57
+
58
+ def _is_cuda_available() -> bool:
59
+ """Check if CUDA is available on the system by trying to run nvidia-smi."""
60
+ try:
61
+ # Attempt to run `nvidia-smi` to check for CUDA.
62
+ # This command should run successfully if NVIDIA drivers are installed and GPUs are present.
63
+ subprocess.run(
64
+ ["nvidia-smi"],
65
+ stdout=subprocess.DEVNULL,
66
+ stderr=subprocess.DEVNULL,
67
+ check=True,
68
+ )
69
+ return True
70
+ except:
71
+ return False
72
+
73
+
74
+ def _handle_device_requests(
75
+ algorithm: AlgorithmData, cuda_devices: str, force_cpu: bool
76
+ ) -> List[docker.types.DeviceRequest]:
77
+ """Handle the device requests for the docker container (request cuda or cpu).
78
+
79
+ Args:
80
+ algorithm (AlgorithmData): Algorithm data
81
+ cuda_devices (str): The CUDA devices to use
82
+ force_cpu (bool): Whether to force CPU execution
83
+ """
84
+ cuda_available = _is_cuda_available()
85
+ if not cuda_available or force_cpu:
86
+ if not algorithm.run_args.cpu_compatible:
87
+ cause = (
88
+ "User tried to force CPU execution but"
89
+ if cuda_available
90
+ else "No Cuda installation/ GPU was found and"
91
+ )
92
+ raise AlgorithmNotCPUCompatibleException(
93
+ f"{cause} the chosen algorithm is not CPU-compatible. Aborting..."
94
+ )
95
+ # empty device requests => run on CPU
96
+ return []
97
+ # request gpu with chosen devices
98
+ return [
99
+ docker.types.DeviceRequest(device_ids=[cuda_devices], capabilities=[["gpu"]])
100
+ ]
101
+
102
+
103
+ def _get_additional_files_path(algorithm: AlgorithmData) -> Path:
104
+ """Get the path to the additional files for @algorithm.
105
+
106
+ Args:
107
+ algorithm (AlgorithmData): The algorithm data
108
+
109
+ Returns:
110
+ Path to the additional files
111
+ """
112
+ # ensure weights are present and get path
113
+ if algorithm.weights is not None:
114
+ return check_model_weights(record_id=algorithm.weights.record_id)
115
+ else:
116
+ # if no weights are directly specified a dummy weights folder will be mounted that is potentially used for parameter files etc.
117
+ return get_dummy_weights_path()
118
+
119
+
120
+ def _get_volume_mappings(
121
+ data_path: Path, additional_files_path: Path, output_path: Path
122
+ ) -> Dict:
123
+ """Get the volume mappings for the docker container.
124
+
125
+ Args:
126
+ data_path (Path): The path to the input data
127
+ additional_files_path (Path): The path to the additional files
128
+ output_path (Path): The path to save the output
129
+
130
+ Returns:
131
+ Dict: The volume mappings
132
+ """
133
+ # TODO: add support for recommended "ro" mount mode for input data
134
+ # data = mlcube_io0, additional files = mlcube_io1, output = mlcube_io2
135
+ return {
136
+ volume.absolute(): {
137
+ "bind": f"/mlcube_io{i}",
138
+ "mode": "rw",
139
+ }
140
+ for i, volume in enumerate([data_path, additional_files_path, output_path])
141
+ }
142
+
143
+
144
+ def _build_args(
145
+ algorithm: AlgorithmData, additional_files_path: Path
146
+ ) -> Tuple[str, str]:
147
+ """Build the command and extra arguments for the docker container.
148
+
149
+ Args:
150
+ algorithm (AlgorithmData): The algorithm data
151
+ additional_files_path (Path): The path to the additional files
152
+
153
+ Returns:
154
+ command_args, extra_args (Tuple): The command arguments and extra arguments
155
+ """
156
+ # Build command that will be run in the docker container
157
+ command_args = (
158
+ f"--data_path=/mlcube_io0 --{algorithm.weights.param_name}=/mlcube_io1 --output_path=/mlcube_io2"
159
+ if algorithm.weights is not None
160
+ else f"--data_path=/mlcube_io0 --output_path=/mlcube_io2"
161
+ )
162
+
163
+ if algorithm.run_args.parameters_file:
164
+ # The algorithms that need a parameters file do not seem to actually use it but just need it to exist
165
+ # As a workaround we simply create an empty file
166
+ parameters_file = additional_files_path / "parameters.yaml"
167
+ parameters_file.touch()
168
+ command_args += f" --parameters_file=/mlcube_io1/parameters.yaml"
169
+
170
+ extra_args = {}
171
+ if not algorithm.run_args.requires_root:
172
+ # run the container as the current user to ensure written files are always owned by the user
173
+ # also overall better security-wise
174
+ extra_args["user"] = f"{os.getuid()}:{os.getgid()}"
175
+
176
+ return command_args, extra_args
177
+
178
+
179
+ def _observe_docker_output(container: docker.models.containers.Container):
180
+ """Observe the output of a running docker container and display a spinner. On Errors log container output.
181
+
182
+ Args:
183
+ container (docker.models.containers.Container): The container to observe
184
+ """
185
+ # capture the output
186
+ container_output = container.attach(
187
+ stdout=True, stderr=True, stream=True, logs=True
188
+ )
189
+
190
+ # Display spinner while the container is running
191
+ with Halo(text="Running inference...", spinner="dots"):
192
+ # Wait for the container to finish
193
+ exit_code = container.wait()
194
+ # Check if the container exited with an error
195
+ if exit_code["StatusCode"] != 0:
196
+ for line in container_output:
197
+ logger.error(f">> {line.decode('utf-8')}")
198
+ raise Exception(
199
+ "Container finished with an error. See logs above for details."
200
+ )
201
+
202
+ # TODO add option to print/ save container output
203
+
204
+
205
+ def run_docker(
206
+ algorithm: AlgorithmData,
207
+ data_path: Path,
208
+ output_path: Path,
209
+ cuda_devices: str,
210
+ force_cpu: bool,
211
+ ):
212
+ """Run a docker container for the provided algorithm.
213
+
214
+ Args:
215
+ algorithm (AlgorithmData): The data of the algorithm to run
216
+ data_path (Path | str): The path to the input data
217
+ output_path (Path | str): The path to save the output
218
+ cuda_devices (str): The CUDA devices to use
219
+ force_cpu (bool): Whether to force CPU execution
220
+ """
221
+ # ensure image is present, if not pull it
222
+ _ensure_image(image=algorithm.run_args.docker_image)
223
+
224
+ # Log the message
225
+ additional_files_path = _get_additional_files_path(algorithm)
226
+
227
+ # ensure output folder exists
228
+ output_path.mkdir(parents=True, exist_ok=True)
229
+
230
+ volume_mappings = _get_volume_mappings(
231
+ data_path=data_path,
232
+ additional_files_path=additional_files_path,
233
+ output_path=output_path,
234
+ )
235
+
236
+ logger.info(f"{'Starting inference'}")
237
+
238
+ command_args, extra_args = _build_args(
239
+ algorithm=algorithm, additional_files_path=additional_files_path
240
+ )
241
+
242
+ # device setup
243
+ device_requests = _handle_device_requests(
244
+ algorithm=algorithm, cuda_devices=cuda_devices, force_cpu=force_cpu
245
+ )
246
+
247
+ # Run the container
248
+ container = client.containers.run(
249
+ image=algorithm.run_args.docker_image,
250
+ volumes=volume_mappings,
251
+ device_requests=device_requests,
252
+ command=f"infer {command_args}",
253
+ network_mode="none",
254
+ detach=True,
255
+ remove=True,
256
+ shm_size=algorithm.run_args.shm_size,
257
+ **extra_args,
258
+ )
259
+ _observe_docker_output(container=container)
260
+
261
+ logger.info(f"{'Finished inference'}")
@@ -0,0 +1,15 @@
1
+ from dacite import DaciteError
2
+
3
+
4
+ class AlgorithmNotCPUCompatibleException(Exception):
5
+ """Exception raised when an CPU-incompatible algorithm tries to run on CPU."""
6
+
7
+ def __init__(self, message: str):
8
+ super().__init__(message)
9
+
10
+
11
+ class AlgorithmConfigException(DaciteError):
12
+ """Exception raised when the algorithm config file has issues."""
13
+
14
+ def __init__(self, message: str):
15
+ super().__init__(message)
@@ -0,0 +1,78 @@
1
+ from __future__ import annotations
2
+
3
+ import shutil
4
+ import sys
5
+ from pathlib import Path
6
+ from typing import Dict, List
7
+
8
+ from loguru import logger
9
+
10
+
11
+ def standardize_subject_inputs(
12
+ data_folder: Path,
13
+ subject_id: str,
14
+ t1c: Path | str,
15
+ t1n: Path | str,
16
+ t2f: Path | str,
17
+ t2w: Path | str,
18
+ ):
19
+ """Standardize the input images for a single subject to match requirements of all algorithms and save them in @data_folder/@internal_name.
20
+ Meaning, e.g. for adult glioma:
21
+ BraTS-GLI-00000-000 \n
22
+ ┣ BraTS-GLI-00000-000-t1c.nii.gz \n
23
+ ┣ BraTS-GLI-00000-000-t1n.nii.gz \n
24
+ ┣ BraTS-GLI-00000-000-t2f.nii.gz \n
25
+ ┗ BraTS-GLI-00000-000-t2w.nii.gz \n
26
+
27
+ Args:
28
+ data_folder (Path): Parent folder where the subject folder will be created
29
+ subject_id (str): Subject ID to be used for the folder and filenames
30
+ t1c (Path | str): T1c image path
31
+ t1n (Path | str): T1n image path
32
+ t2f (Path | str): T2f image path
33
+ t2w (Path | str): T2w image path
34
+ """
35
+ subject_folder = data_folder / subject_id
36
+ subject_folder.mkdir(parents=True, exist_ok=True)
37
+ # TODO: investigate usage of symlinks (might cause issues on windows and would probably require different volume handling)
38
+ try:
39
+ shutil.copy(t1c, subject_folder / f"{subject_id}-t1c.nii.gz")
40
+ shutil.copy(t1n, subject_folder / f"{subject_id}-t1n.nii.gz")
41
+ shutil.copy(t2f, subject_folder / f"{subject_id}-t2f.nii.gz")
42
+ shutil.copy(t2w, subject_folder / f"{subject_id}-t2w.nii.gz")
43
+ except FileNotFoundError as e:
44
+ logger.error(f"Error while standardizing files: {e}")
45
+ logger.error(
46
+ "Please make sure the input files are in the correct format, i.e.:\n A/A-t1c.nii.gz, A/A-t1n.nii.gz, A/A-t2f.nii.gz, A/A-t2w.nii.gz"
47
+ )
48
+ sys.exit(1)
49
+
50
+
51
+ def standardize_subjects_inputs_list(
52
+ subjects: List[Path], temp_data_folder: Path, input_name_schema: str
53
+ ) -> Dict[str, str]:
54
+ """Standardize the input images for a list of subjects to match requirements of all algorithms and save them in @temp_data_folder/@subject_id.
55
+
56
+ Args:
57
+ subjects (List[Path]): List of subject folders, each with a t1c, t1n, t2f, t2w image in standard format
58
+ temp_data_folder (Path): Parent folder where the subject folders will be created
59
+ input_name_schema (str): Schema to be used for the subject folder and filenames depending on the BraTS Challenge
60
+
61
+ Returns:
62
+ Dict[str, str]: Dictionary mapping internal name (in standardized format) to external subject name provided by user
63
+ """
64
+ internal_external_name_map = {}
65
+ for i, subject in enumerate(subjects):
66
+ internal_name = input_name_schema.format(id=i)
67
+ internal_external_name_map[internal_name] = subject.name
68
+ # TODO Add support for .nii files
69
+
70
+ standardize_subject_inputs(
71
+ data_folder=temp_data_folder,
72
+ subject_id=internal_name,
73
+ t1c=subject / f"{subject.name}-t1c.nii.gz",
74
+ t1n=subject / f"{subject.name}-t1n.nii.gz",
75
+ t2f=subject / f"{subject.name}-t2f.nii.gz",
76
+ t2w=subject / f"{subject.name}-t2w.nii.gz",
77
+ )
78
+ return internal_external_name_map
@@ -0,0 +1,179 @@
1
+ from __future__ import annotations
2
+
3
+ import shutil
4
+ import sys
5
+ import zipfile
6
+ from io import BytesIO
7
+ from pathlib import Path
8
+ from typing import Dict, List
9
+
10
+ import requests
11
+ from loguru import logger
12
+ from tqdm import tqdm
13
+
14
+ ZENODO_RECORD_BASE_URL = "https://zenodo.org/api/records"
15
+ WEIGHTS_FOLDER = Path(__file__).parent / "weights"
16
+
17
+
18
+ def get_dummy_weights_path() -> Path:
19
+ dummy = WEIGHTS_FOLDER / "dummy"
20
+ dummy.mkdir(exist_ok=True, parents=True)
21
+ return dummy
22
+
23
+
24
+ def check_model_weights(record_id: str) -> Path:
25
+ """Check if latest model weights are present and download them otherwise.
26
+
27
+ Returns:
28
+ Path: Path to the model weights folder.
29
+ """
30
+
31
+ zenodo_metadata, archive_url = _get_zenodo_metadata_and_archive_url(
32
+ record_id=record_id
33
+ )
34
+
35
+ record_weights_pattern = f"{record_id}_v*.*.*"
36
+ matching_folders = list(WEIGHTS_FOLDER.glob(record_weights_pattern))
37
+ # Get the latest downloaded weights
38
+ latest_downloaded_weights = _get_latest_version_folder_name(matching_folders)
39
+
40
+ if not latest_downloaded_weights:
41
+ if not zenodo_metadata:
42
+ logger.error(
43
+ "Model weights not found locally and Zenodo could not be reached. Exiting..."
44
+ )
45
+ sys.exit()
46
+ logger.info(f"Model weights not found locally")
47
+
48
+ return _download_model_weights(
49
+ zenodo_metadata=zenodo_metadata,
50
+ record_id=record_id,
51
+ archive_url=archive_url,
52
+ )
53
+
54
+ logger.info(f"Found downloaded local weights: {latest_downloaded_weights}")
55
+
56
+ if not zenodo_metadata:
57
+ logger.warning(
58
+ "Zenodo server could not be reached. Using the latest downloaded weights."
59
+ )
60
+ return WEIGHTS_FOLDER / latest_downloaded_weights
61
+
62
+ # Compare the latest downloaded weights with the latest Zenodo version
63
+ if zenodo_metadata["version"] == latest_downloaded_weights.split("_v")[1]:
64
+ logger.info(
65
+ f"Latest model weights ({latest_downloaded_weights}) are already present."
66
+ )
67
+ return WEIGHTS_FOLDER / latest_downloaded_weights
68
+
69
+ logger.info(
70
+ f"New model weights available on Zenodo ({zenodo_metadata['version']}). Deleting old and fetching new weights..."
71
+ )
72
+ # delete old weights
73
+ shutil.rmtree(
74
+ WEIGHTS_FOLDER / latest_downloaded_weights,
75
+ onerror=lambda func, path, excinfo: logger.warning(
76
+ f"Failed to delete {path}: {excinfo}"
77
+ ),
78
+ )
79
+ return _download_model_weights(
80
+ zenodo_metadata=zenodo_metadata, record_id=record_id, archive_url=archive_url
81
+ )
82
+
83
+
84
+ def _get_latest_version_folder_name(folders: List[Path]) -> str | None:
85
+ """Get the latest (non empty) version folder name from the list of folders.
86
+
87
+ Args:
88
+ folders (List[Path]): List of folders matching the pattern.
89
+
90
+ Returns:
91
+ str | None: Latest version folder name if one exists, else None.
92
+ """
93
+ if not folders:
94
+ return None
95
+ latest_downloaded_folder = sorted(
96
+ folders,
97
+ reverse=True,
98
+ key=lambda x: tuple(map(int, str(x).split("_v")[1].split("."))),
99
+ )[0]
100
+ # check folder is not empty
101
+ if not list(latest_downloaded_folder.glob("*")):
102
+ return None
103
+ return latest_downloaded_folder.name
104
+
105
+
106
+ def _get_zenodo_metadata_and_archive_url(record_id: str) -> Dict | None:
107
+ """Get the metadata for the Zenodo record and the files archive url.
108
+
109
+ Returns:
110
+ Tuple: (dict: Metadata for the Zenodo record, str: URL to the archive file)
111
+ """
112
+ try:
113
+ response = requests.get(f"{ZENODO_RECORD_BASE_URL}/{record_id}")
114
+ if response.status_code != 200:
115
+ logger.error(
116
+ f"Cant find model weights for record_id '{record_id}' on Zenodo. Exiting..."
117
+ )
118
+ data = response.json()
119
+ return data["metadata"], data["links"]["archive"]
120
+
121
+ except requests.exceptions.RequestException as e:
122
+ logger.warning(f"Failed to fetch Zenodo metadata: {e}")
123
+ return None
124
+
125
+
126
+ def _download_model_weights(
127
+ zenodo_metadata: Dict, record_id: str, archive_url: str
128
+ ) -> Path:
129
+ """Download the latest model weights from Zenodo for the requested record and extract them to the target folder.
130
+
131
+ Args:
132
+ weights_folder (Path): General weights folder path in which the requested model weights will be stored.
133
+ zenodo_metadata (Dict): Metadata for the Zenodo record.
134
+ record_id (str): Zenodo record ID.
135
+
136
+ Returns:
137
+ Path: Path to the model weights folder for the requested record.
138
+ """
139
+ record_weights_folder = (
140
+ WEIGHTS_FOLDER / f"{record_id}_v{zenodo_metadata['version']}"
141
+ )
142
+ # ensure folder exists
143
+ record_weights_folder.mkdir(parents=True, exist_ok=True)
144
+
145
+ logger.info(f"Downloading model weights from Zenodo. This might take a while...")
146
+ # Make a GET request to the URL
147
+ response = requests.get(archive_url, stream=True)
148
+ # Ensure the request was successful
149
+ if response.status_code != 200:
150
+ logger.error(
151
+ f"Failed to download model weights. Status code: {response.status_code}"
152
+ )
153
+ return
154
+
155
+ # Download with progress bar
156
+ chunk_size = 1024 # 1KB
157
+ bytes_io = BytesIO()
158
+ with tqdm(
159
+ total=0, # unknown size since content length not given
160
+ unit="B",
161
+ unit_scale=True,
162
+ ) as pbar:
163
+ for data in response.iter_content(chunk_size=chunk_size):
164
+ bytes_io.write(data)
165
+ pbar.update(len(data))
166
+
167
+ # Extract the downloaded zip file to the target folder
168
+ with zipfile.ZipFile(bytes_io) as zip_ref:
169
+ zip_ref.extractall(record_weights_folder)
170
+
171
+ # check if the extracted file is still a zip
172
+ for f in record_weights_folder.iterdir():
173
+ if f.is_file() and f.suffix == ".zip":
174
+ with zipfile.ZipFile(f) as zip_ref:
175
+ zip_ref.extractall(record_weights_folder)
176
+ f.unlink() # remove zip after extraction
177
+
178
+ logger.info(f"Zip file extracted successfully to {record_weights_folder}")
179
+ return record_weights_folder
@@ -0,0 +1,18 @@
1
+ [project]
2
+ name = "BraTS"
3
+ version = "0.0.1"
4
+ authors = [
5
+ { name="Example Author", email="author@example.com" },
6
+ ]
7
+ description = "A small example package"
8
+ readme = "README.md"
9
+ requires-python = ">=3.8"
10
+ classifiers = [
11
+ "Programming Language :: Python :: 3",
12
+ "License :: OSI Approved :: MIT License",
13
+ "Operating System :: OS Independent",
14
+ ]
15
+
16
+ [project.urls]
17
+ Homepage = "https://github.com/pypa/sampleproject"
18
+ Issues = "https://github.com/pypa/sampleproject/issues"
brats-0.0.1/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,20 @@
1
+ import unittest
2
+ from pathlib import Path
3
+
4
+ from brats.algorithm_config import load_algorithms
5
+
6
+ PACKAGE_DIR = Path(__file__).parent.parent
7
+
8
+
9
+ class TestAlgorithmConfig(unittest.TestCase):
10
+
11
+ def test_configs_valid(self):
12
+ algorithms_folder = PACKAGE_DIR / "brats" / "algorithms"
13
+
14
+ configs = algorithms_folder.iterdir()
15
+
16
+ for config in configs:
17
+ try:
18
+ load_algorithms(file_path=config)
19
+ except Exception as e:
20
+ self.fail(f"Failed to load config {config}: {e}")
@@ -0,0 +1,84 @@
1
+ import shutil
2
+ import tempfile
3
+ import unittest
4
+ from pathlib import Path
5
+ from unittest.mock import patch
6
+
7
+ from loguru import logger
8
+
9
+ from brats.utils import standardize_subject_inputs, standardize_subjects_inputs_list
10
+
11
+
12
+ class TestStandardizeInputs(unittest.TestCase):
13
+ def setUp(self):
14
+ # Create a temporary directory
15
+ self.test_dir = Path(tempfile.mkdtemp())
16
+ self.data_folder = self.test_dir / "data"
17
+ self.data_folder.mkdir(parents=True, exist_ok=True)
18
+ self.temp_data_folder = self.test_dir / "tmp_std_data"
19
+ self.temp_data_folder.mkdir(parents=True, exist_ok=True)
20
+
21
+ # Create mock paths for input images
22
+ self.subject_folder = self.data_folder / "subject"
23
+ self.subject_folder.mkdir(parents=True, exist_ok=True)
24
+ self.t1c = self.data_folder / "subject-t1c.nii.gz"
25
+ self.t1n = self.data_folder / "subject-t1n.nii.gz"
26
+ self.t2f = self.data_folder / "subject-t2f.nii.gz"
27
+ self.t2w = self.data_folder / "subject-t2w.nii.gz"
28
+ # Create dummy files
29
+ for img in [self.t1c, self.t1n, self.t2f, self.t2w]:
30
+ img.touch(exist_ok=True)
31
+
32
+ def tearDown(self):
33
+ # Remove the temporary directory after the test
34
+ shutil.rmtree(self.test_dir)
35
+
36
+ def test_successful_standardization(self):
37
+ subject_id = "test_subject"
38
+ standardize_subject_inputs(
39
+ data_folder=self.temp_data_folder,
40
+ subject_id=subject_id,
41
+ t1c=self.t1c,
42
+ t1n=self.t1n,
43
+ t2f=self.t2f,
44
+ t2w=self.t2w,
45
+ )
46
+ subject_folder = self.temp_data_folder / subject_id
47
+ self.assertTrue(subject_folder.exists())
48
+ for img_type in ["t1c", "t1n", "t2f", "t2w"]:
49
+ self.assertTrue(
50
+ (subject_folder / f"{subject_id}-{img_type}.nii.gz").exists()
51
+ )
52
+
53
+ @patch("sys.exit")
54
+ @patch.object(logger, "error")
55
+ def test_handle_file_not_found_error(self, mock_logger, mock_exit):
56
+ subject_id = "test_subject"
57
+ # Provide a non-existent file path for t1c
58
+ t1c = "non_existent_file.nii.gz"
59
+ standardize_subject_inputs(
60
+ data_folder=self.data_folder,
61
+ subject_id=subject_id,
62
+ t1c=t1c,
63
+ t1n=self.t1n,
64
+ t2f=self.t2f,
65
+ t2w=self.t2w,
66
+ )
67
+ mock_logger.assert_called()
68
+ mock_exit.assert_called_with(1)
69
+
70
+ @patch("brats.utils.standardize_subject_inputs")
71
+ def test_standardize_subjects_inputs_list(self, mock_standardize_subject_inputs):
72
+ subjects = [f for f in self.data_folder.iterdir() if f.is_dir()]
73
+ mapping = standardize_subjects_inputs_list(
74
+ subjects=subjects,
75
+ temp_data_folder=self.temp_data_folder,
76
+ input_name_schema="BraTS-PED-{id:05d}-000",
77
+ )
78
+ self.assertDictEqual(
79
+ mapping,
80
+ {
81
+ "BraTS-PED-00000-000": "subject",
82
+ },
83
+ )
84
+ mock_standardize_subject_inputs.assert_called_once()