diff --git a/README.md b/README.md index 5cd4a73..19fab74 100644 --- a/README.md +++ b/README.md @@ -25,12 +25,12 @@ pip install brats ### Docker and NVIDIA Container Toolkit Setup To run `brats` you need a working Docker installation. -Most algorithms also require GPU support (NVIDIA Docker). +Most algorithms also require GPU support (NVIDIA Docker). Installation instructions: -- **Docker**: Installation instructions on the official [website](https://docs.docker.com/get-docker/) -- **NVIDIA Container Toolkit**: Refer to the [NVIDIA install guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) and the official [GitHub page](https://github.com/NVIDIA/nvidia-container-toolkit) +- **Docker**: Installation instructions on the official [website](https://docs.docker.com/get-docker/) +- **NVIDIA Container Toolkit**: Refer to the [NVIDIA install guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) and the official [GitHub page](https://github.com/NVIDIA/nvidia-container-toolkit) ## Use Cases and Tutorials @@ -51,32 +51,72 @@ segmenter.infer_single( For more examples and details please refer to our extensive Notebook tutorials [**TODO**] ## Algorithms -| Challenge | Year | Position | Authors | Paper | CPU support | -|---|---|---|---|---|---| -| *Meningioma Segmentation* | 2023 | 1st | Andriy Myronenko, et al. | N/A | :heavy_multiplication_x: | -| | 2023 | 2nd | Ziyan Huang | N/A | :white_check_mark: | -| | 2023 | 3rd | Zhifan Jiang et al. | N/A | :heavy_multiplication_x: | -| *Pediatric Segmentation* | 2023 | 1st | Zhifan Jiang et al. | N/A | :heavy_multiplication_x: | -| | 2023 | 2nd | Andriy Myronenko, et al. | N/A | :heavy_multiplication_x: | -| | 2023 | 3rd | Yubo Zhou | N/A | :heavy_multiplication_x: | -| *Adult Glioma Segmentation* | 2023 | 1st | André Ferreira, et al. | [Link](https://arxiv.org/abs/2402.17317v1) | :heavy_multiplication_x: | -| | 2023 | 2nd | Andriy Myronenko, et al. | N/A | :heavy_multiplication_x: | -| | 2023 | 3rd | Fadillah Adamsyah Maani, et al. | N/A | :heavy_multiplication_x: | + +
+ + Adult Glioma Segmentation +
+ +| Year | Rank | Author | Paper | CPU Support | +| ---- | ---- | --------------------------------- | ------------------------------------------ | ----------- | +| 2023 | 1st | _André Ferreira, et al._ | [Link](https://arxiv.org/abs/2402.17317v1) | ❌ | +| 2023 | 2nd | _Andriy Myronenko, et al._ | N/A | ❌ | +| 2023 | 3rd | _Fadillah Adamsyah Maani, et al._ | N/A | ❌ | + +
+ +
+ BraTS-Africa Segmentation +
+ +| Year | Rank | Author | Paper | CPU Support | +| ---- | ---- | -------------------------- | ----- | ----------- | +| 2023 | 1st | _Andriy Myronenko, et al._ | N/A | ❌ | +| 2023 | 2nd | _Alyssa R Amod, et al._ | N/A | ❌ | +| 2023 | 3rd | _Ziyan Huang, et al._ | N/A | ✅ | + +
+ +
+ Meningioma Segmentation +
+ +| Year | Rank | Author | Paper | CPU Support | +| ---- | ---- | -------------------------- | ----- | ----------- | +| 2023 | 1st | _Andriy Myronenko, et al._ | N/A | ❌ | +| 2023 | 2nd | _Ziyan Huang_ | N/A | ✅ | +| 2023 | 3rd | _Zhifan Jiang et al._ | N/A | ❌ | + +
+ +
+ Pediatric Segmentation +
+ +| Year | Rank | Author | Paper | CPU Support | +| ---- | ---- | -------------------------- | ----- | ----------- | +| 2023 | 1st | _Zhifan Jiang et al._ | N/A | ❌ | +| 2023 | 2nd | _Andriy Myronenko, et al._ | N/A | ❌ | +| 2023 | 3rd | _Yubo Zhou_ | N/A | ❌ | + +
## Citation + If you use BraTS in your research, please cite it to support the development! ``` TODO: citation will be added asap ``` - ## Contributing We welcome all kinds of contributions from the community! ### Reporting Bugs, Feature Requests and Questions + Please open a new issue [here](https://github.com/BrainLesion/BraTS/issues). ### Code contributions + Nice to have you on board! Please have a look at our [CONTRIBUTING.md](CONTRIBUTING.md) file. diff --git a/brats/__init__.py b/brats/__init__.py index 0b4faf3..a1de91c 100644 --- a/brats/__init__.py +++ b/brats/__init__.py @@ -2,4 +2,5 @@ AdultGliomaSegmenter, MeningiomaSegmenter, PediatricSegmenter, + AfricaSegmenter, ) diff --git a/brats/algorithm_config.py b/brats/algorithm_config.py index e1a7704..be15835 100644 --- a/brats/algorithm_config.py +++ b/brats/algorithm_config.py @@ -17,6 +17,10 @@ class MetaData: """If available, a url to the paper of the algorithm""" challenge: str """The challenge the algorithm was submitted to""" + rank: str + """The rank of the algorithm in the challenge""" + year: int + """The year the algorithm was submitted""" @dataclass diff --git a/brats/algorithms.py b/brats/algorithms.py index 00bfca4..8a0283e 100644 --- a/brats/algorithms.py +++ b/brats/algorithms.py @@ -11,10 +11,12 @@ from brats.algorithm_config import load_algorithms from brats.constants import ( ADULT_GLIOMA_SEGMENTATION_ALGORITHMS, + AFRICA_SEGMENTATION_ALGORITHMS, MENINGIOMA_SEGMENTATION_ALGORITHMS, PEDIATRIC_SEGMENTATION_ALGORITHMS, AdultGliomaAlgorithms, Algorithms, + AfricaAlgorithms, MeningiomaAlgorithms, PediatricAlgorithms, ) @@ -56,7 +58,7 @@ def __init__( def _log_algorithm_info(self): """Log information about the selected algorithm.""" logger.opt(colors=True).info( - f"Running algorithm: {self.algorithm.meta.challenge}" + f"Running algorithm: {self.algorithm.meta.challenge} [{self.algorithm.meta.rank} place]" ) logger.opt(colors=True).info( f"(Paper) Consider citing the corresponding paper: {self.algorithm.meta.paper} by {self.algorithm.meta.authors}" @@ -80,6 +82,24 @@ def _add_log_file_handler(self, log_file: Path | str) -> int: return logger_id + def _cleanup( + self, temp_data_folder: Path, temp_output_folder: Path, logger_id: Optional[int] + ) -> None: + try: + shutil.rmtree(temp_data_folder) + except PermissionError as e: + logger.warning( + f"Failed to remove temporary folder {temp_data_folder}. This is most likely caused by bad permission management of the docker container. \nError: {e}" + ) + try: + shutil.rmtree(temp_output_folder) + except PermissionError as e: + logger.warning( + f"Failed to remove temporary folder {temp_output_folder}. This is most likely caused by bad permission management of the docker container. \nError: {e}" + ) + if logger_id is not None: + logger.remove(logger_id) + def infer_single( self, t1c: Path | str, @@ -137,10 +157,11 @@ def infer_single( logger.info(f"Saved segmentation to: {output_file.absolute()}") finally: - shutil.rmtree(temp_data_folder) - shutil.rmtree(temp_output_folder) - if log_file is not None: - logger.remove(logger_id) + self._cleanup( + temp_data_folder=temp_data_folder, + temp_output_folder=temp_output_folder, + logger_id=logger_id if log_file else None, + ) def infer_batch( self, @@ -207,10 +228,11 @@ def infer_batch( logger.info(f"Saved results to: {output_folder.absolute()}") finally: - shutil.rmtree(temp_data_folder) - shutil.rmtree(temp_output_folder) - if log_file: - logger.remove(logger_id) + self._cleanup( + temp_data_folder=temp_data_folder, + temp_output_folder=temp_output_folder, + logger_id=logger_id if log_file else None, + ) class AdultGliomaSegmenter(BraTSAlgorithm): @@ -240,7 +262,7 @@ class MeningiomaSegmenter(BraTSAlgorithm): """Provides algorithms to perform tumor segmentation on adult meningioma MRI data. Args: - algorithm (AdultGliomaAlgorithms, optional): Select an algorithm. Defaults to MeningiomaAlgorithms.BraTS23_1. + algorithm (MeningiomaAlgorithms, optional): Select an algorithm. Defaults to MeningiomaAlgorithms.BraTS23_1. cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0". force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False. """ @@ -263,7 +285,7 @@ class PediatricSegmenter(BraTSAlgorithm): """Provides algorithms to perform tumor segmentation on pediatric MRI data Args: - algorithm (AdultGliomaAlgorithms, optional): Select an algorithm. Defaults to PediatricAlgorithms.BraTS23_1. + algorithm (PediatricAlgorithms, optional): Select an algorithm. Defaults to PediatricAlgorithms.BraTS23_1. cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0". force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False. """ @@ -280,3 +302,26 @@ def __init__( cuda_devices=cuda_devices, force_cpu=force_cpu, ) + + +class AfricaSegmenter(BraTSAlgorithm): + """Provides algorithms to perform tumor segmentation on data from the BraTSAfrica challenge + + Args: + algorithm (AfricaAlgorithms, optional): Select an algorithm. Defaults to AfricaAlgorithms.BraTS23_1. + cuda_devices (Optional[str], optional): Which cuda devices to use. Defaults to "0". + force_cpu (bool, optional): Execution will default to GPU, this flag allows forced CPU execution if the algorithm is compatible. Defaults to False. + """ + + def __init__( + self, + algorithm: AfricaAlgorithms = AfricaAlgorithms.BraTS23_1, + cuda_devices: Optional[str] = "0", + force_cpu: bool = False, + ): + super().__init__( + algorithm=algorithm, + algorithms_file_path=AFRICA_SEGMENTATION_ALGORITHMS, + cuda_devices=cuda_devices, + force_cpu=force_cpu, + ) diff --git a/brats/algorithms/adult_glioma.yml b/brats/algorithms/adult_glioma.yml index d928e45..b52eb74 100644 --- a/brats/algorithms/adult_glioma.yml +++ b/brats/algorithms/adult_glioma.yml @@ -4,7 +4,9 @@ algorithms: meta: authors: André Ferreira, et al. paper: https://arxiv.org/abs/2402.17317v1 - challenge: BraTS23 Adult Glioma Segmentation (1st place) + challenge: BraTS23 Adult Glioma Segmentation + rank: 1st + year: 2023 run_args: docker_image: brainles/brats23_faking_it:latest input_name_schema: "BraTS-GLI-{id:05d}-000" @@ -16,7 +18,9 @@ algorithms: meta: authors: Andriy Myronenko, et al. paper: N/A - challenge: BraTS23 Adult Glioma Segmentation (2nd place) + challenge: BraTS23 Adult Glioma Segmentation + rank: 2nd + year: 2023 run_args: docker_image: brainles/brats23_nvauto:latest input_name_schema: "BraTS-GLI-{id:05d}-000" @@ -28,7 +32,9 @@ algorithms: meta: authors: Fadillah Adamsyah Maani, et al. paper: N/A - challenge: BraTS23 Adult Glioma Segmentation (3rd place) + challenge: BraTS23 Adult Glioma Segmentation + rank: 3rd + year: 2023 run_args: docker_image: brainles/brats23_biomedmbz:latest input_name_schema: "BraTS-GLI-{id:05d}-000" diff --git a/brats/algorithms/africa.yml b/brats/algorithms/africa.yml new file mode 100644 index 0000000..626ec6a --- /dev/null +++ b/brats/algorithms/africa.yml @@ -0,0 +1,45 @@ +algorithms: + + BraTS23_1: + meta: + authors: Andriy Myronenko, et al. + paper: TODO + challenge: BraTS23 BraTS-Africa Segmentation + rank: 1st + year: 2023 + run_args: + docker_image: brainles/brats23_africa_nvauto:latest + input_name_schema: "BraTS-SSA-{id:05d}-000" + requires_root: true + parameters_file: true + shm_size: "32gb" + + BraTS23_2: + meta: + authors: Alyssa R Amod, et al. + paper: N/A + challenge: BraTS23 BraTS-Africa Segmentation + rank: 2nd + year: 2023 + run_args: + docker_image: brainles/brats23_africa_sparkunn:latest + input_name_schema: "BraTS-SSA-{id:05d}-000" + requires_root: false + parameters_file: true + weights: + record_id: "13373752" + param_name: "ckpts_path" + BraTS23_3: + meta: + authors: Ziyan Huang, et al. + paper: N/A + challenge: BraTS23 BraTS-Africa Segmentation + rank: 3rd + year: 2023 + run_args: + docker_image: brainles/brats23_africa_blackbean:latest + input_name_schema: "BraTS-SSA-{id:05d}-000" + requires_root: true + parameters_file: true + cpu_compatible: true + diff --git a/brats/algorithms/meningioma.yml b/brats/algorithms/meningioma.yml index db05f35..2deea67 100644 --- a/brats/algorithms/meningioma.yml +++ b/brats/algorithms/meningioma.yml @@ -3,7 +3,9 @@ algorithms: meta: authors: Andriy Myronenko, et al. paper: N/A - challenge: BraTS23 Meningioma Segmentation (1st place) + challenge: BraTS23 Meningioma Segmentation + rank: 1st + year: 2023 run_args: docker_image: brainles/brats23_meningioma_nvauto:latest input_name_schema: "BraTS-MEN-{id:05d}-000" @@ -15,7 +17,9 @@ algorithms: meta: authors: Ziyan Huang paper: N/A - challenge: BraTS23 Meningioma Segmentation (2nd place) + challenge: BraTS23 Meningioma Segmentation + rank: 2nd + year: 2023 run_args: docker_image: brainles/brats23_meningioma_blackbean:latest input_name_schema: "BraTS-MEN-{id:05d}-000" @@ -28,7 +32,9 @@ algorithms: meta: authors: Zhifan Jiang et al. paper: N/A - challenge: BraTS23 Meningioma Segmentation (3rd place) + challenge: BraTS23 Meningioma Segmentation + rank: 3rd + year: 2023 run_args: docker_image: brainles/brats23_meningioma_cnmc_pmi2023:latest input_name_schema: "BraTS-MEN-{id:05d}-000" diff --git a/brats/algorithms/parameters/brats23_africa_sparkunn.yml b/brats/algorithms/parameters/brats23_africa_sparkunn.yml new file mode 100644 index 0000000..353f11d --- /dev/null +++ b/brats/algorithms/parameters/brats23_africa_sparkunn.yml @@ -0,0 +1,63 @@ +# Arguments of the data_preparation.py script + +# Arguments of the preprocess.py script +prep_exec_mode: test +ohe: true +verbose: true +task: "12" # testing +dim: 3 +n_jobs: -1 + +# Arguments of the main.py script +exec_mode: predict +config: None +logname : logs.json +gpus: 1 +nodes: 1 +learning_rate: 0.0003 +gradient_clip_val: 0 +negative_slope: 0.01 +tta: true +brats: true +deep_supervision: false +invert_resampled_y: false +amp: true +benchmark: false +focal: false +save_ckpt: false +nfolds: 10 +seed: null # "Random seed +skip_first_n_eval: 0 # Skip the evaluation for the first n epochs. +# ckpt_path: null # Path for loading checkpoint #? +# ckpt_store_dir: $results # Path for saving checkpoint #? +# fold: 2 # Fold number +patience: 100 # Early stopping patience +batch_size: 2 # Batch size +val_batch_size: 4 # Validation batch size +momentum: 0.99 # Momentum factor +weight_decay: 0.0001 #Weight decay (L2 penalty) +save_preds: true # Enable prediction saving +dim: 3 # UNet dimension +resume_training: false # Resume training from the last checkpoint +num_workers: 8 # Number of subprocesses to use for data loading +epochs: 150 # Number of training epochs. +warmup: 5 # Warmup iterations before collecting statistics +nvol: 4 #Number of volumes which come into single batch size for 2D model +depth: 5 # The depth of the encoder +min_fmap: 4 #Minimal dimension of feature map in the bottleneck +deep_supr_num: 2 # Number of deep supervision heads +res_block: false # Enable residual blocks +filters: null # nargs="+", help="[Optional] Set U-Net filters" +layout: NCDHW +brats22_model: true # Use BraTS22 model +norm: instance # Normalization layer +data2d_dim: 3 # Input data dimension for 2d model +oversampling: 0.4 # Probability of crop to have some region with positive label +overlap: 0.25 # Amount of overlap between scans during sliding window inference +scheduler: true # Enable cosine rate scheduler with warmup +optimizer: adam # Optimizer +blend: constant # How to blend output of overlapping windows +train_batches: 0 # Limit number of batches for training (used for benchmarking mode only) +test_batches: 0 # Limit number of batches for inference (used for benchmarking mode only) + +# Argument for the postprocess.py script \ No newline at end of file diff --git a/brats/algorithms/parameters/dummy.yml b/brats/algorithms/parameters/dummy.yml new file mode 100644 index 0000000..e69de29 diff --git a/brats/algorithms/pediatric.yml b/brats/algorithms/pediatric.yml index a3c0320..df8d194 100644 --- a/brats/algorithms/pediatric.yml +++ b/brats/algorithms/pediatric.yml @@ -5,7 +5,9 @@ algorithms: meta: authors: Zhifan Jiang et al. paper: N/A - challenge: BraTS23 pediatric Segmentation (1st place) + challenge: BraTS23 Pediatric Segmentation + rank: 1st + year: 2023 run_args: docker_image: brainles/brats23_pediatric_cnmc_pmi2023:latest input_name_schema: "BraTS-PED-{id:05d}-000" @@ -17,7 +19,9 @@ algorithms: meta: authors: Andriy Myronenko, et al. paper: N/A - challenge: BraTS23 pediatric Segmentation (2nd place) + challenge: BraTS23 Pediatric Segmentation + rank: 2nd + year: 2023 run_args: docker_image: brainles/brats23_pediatric_nvauto:latest input_name_schema: "BraTS-PED-{id:05d}-000" @@ -29,7 +33,9 @@ algorithms: meta: authors: Yubo Zhou paper: N/A - challenge: BraTS23 pediatric Segmentation (3rd place) + challenge: BraTS23 Pediatric Segmentation + rank: 3rd + year: 2023 run_args: docker_image: brainles/brats23_pediatric_sherlock_zyb:latest input_name_schema: "BraTS-PED-{id:05d}-000" diff --git a/brats/constants.py b/brats/constants.py index 6d19200..c7ce4fb 100644 --- a/brats/constants.py +++ b/brats/constants.py @@ -41,8 +41,25 @@ class PediatricAlgorithms(Algorithms): """BraTS23 Pediatric Segmentation 3rd place (GPU only)""" +class AfricaAlgorithms(Algorithms): + """Constants for the available africa segmentation algorithms.""" + + 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)""" + + # meta data file paths ALGORITHM_DIR = Path(__file__).parent / "algorithms" +PARAMETERS_DIR = ALGORITHM_DIR / "parameters" +DUMMY_PARAMETERS = PARAMETERS_DIR / "dummy.yml" ADULT_GLIOMA_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "adult_glioma.yml" MENINGIOMA_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "meningioma.yml" PEDIATRIC_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "pediatric.yml" +AFRICA_SEGMENTATION_ALGORITHMS = ALGORITHM_DIR / "africa.yml" + +WEIGHTS_FOLDER = Path(__file__).parent / "weights" +ZENODO_RECORD_BASE_URL = "https://zenodo.org/api/records" diff --git a/brats/docker.py b/brats/docker.py index 4ab25cd..e34e2dc 100644 --- a/brats/docker.py +++ b/brats/docker.py @@ -1,20 +1,22 @@ from __future__ import annotations import os +import shutil import subprocess -from pathlib import Path -from typing import Dict, List, Tuple import time +from pathlib import Path +from typing import Dict, List, Optional, Tuple import docker +from docker.errors import DockerException from loguru import logger -from rich.progress import Progress from rich.console import Console +from rich.progress import Progress from brats.algorithm_config import AlgorithmData +from brats.constants import PARAMETERS_DIR, DUMMY_PARAMETERS from brats.exceptions import AlgorithmNotCPUCompatibleException, BraTSContainerException from brats.weights import check_model_weights, get_dummy_weights_path -from docker.errors import DockerException try: client = docker.from_env() @@ -97,10 +99,12 @@ def _handle_device_requests( if cuda_available else "No Cuda installation/ GPU was found and" ) + # TODO add reference to table of cpu capable algos as help! raise AlgorithmNotCPUCompatibleException( f"{cause} the chosen algorithm is not CPU-compatible. Aborting..." ) # empty device requests => run on CPU + logger.info("Forcing CPU execution") return [] # request gpu with chosen devices return [ @@ -126,7 +130,10 @@ def _get_additional_files_path(algorithm: AlgorithmData) -> Path: def _get_volume_mappings( - data_path: Path, additional_files_path: Path, output_path: Path + data_path: Path, + additional_files_path: Path, + output_path: Path, + parameters_path: Path, ) -> Dict: """Get the volume mappings for the docker container. @@ -134,29 +141,51 @@ def _get_volume_mappings( data_path (Path): The path to the input data additional_files_path (Path): The path to the additional files output_path (Path): The path to save the output + parameters_path (Path): The path to mount for the parameters file Returns: Dict: The volume mappings """ # TODO: add support for recommended "ro" mount mode for input data - # data = mlcube_io0, additional files = mlcube_io1, output = mlcube_io2 + # data = mlcube_io0, additional files = mlcube_io1, output = mlcube_io2, parameters = mlcube_io3 return { volume.absolute(): { "bind": f"/mlcube_io{i}", "mode": "rw", } - for i, volume in enumerate([data_path, additional_files_path, output_path]) + for i, volume in enumerate( + [data_path, additional_files_path, output_path, parameters_path] + ) } +def _get_parameters_arg(algorithm: AlgorithmData) -> Optional[str]: + """Get the parameters argument for the docker container. + + Args: + algorithm (AlgorithmData): The algorithm data + + Returns: + Optional[str]: The parameters argument for the docker container or None if a parameter file is not required + """ + if algorithm.run_args.parameters_file: + # Docker image name is used as the identifier for the param file + identifier = algorithm.run_args.docker_image.split(":")[0].split("/")[-1] + file = PARAMETERS_DIR / f"{identifier}.yml" + # Some algorithms do require a param file to be present but don't actually use it + # In this case we simply use a dummy file + param_file = file if file.exists() else DUMMY_PARAMETERS + return f" --parameters_file=/mlcube_io3/{param_file.name}" + return None + + def _build_args( - algorithm: AlgorithmData, additional_files_path: Path + algorithm: AlgorithmData, ) -> Tuple[str, str]: """Build the command and extra arguments for the docker container. Args: algorithm (AlgorithmData): The algorithm data - additional_files_path (Path): The path to the additional files Returns: command_args, extra_args (Tuple): The command arguments and extra arguments @@ -168,12 +197,10 @@ def _build_args( else f"--data_path=/mlcube_io0 --output_path=/mlcube_io2" ) - if algorithm.run_args.parameters_file: - # The algorithms that need a parameters file do not seem to actually use it but just need it to exist - # As a workaround we simply create an empty file - parameters_file = additional_files_path / "parameters.yaml" - parameters_file.touch() - command_args += f" --parameters_file=/mlcube_io1/parameters.yaml" + # Add parameters file arg if required + params_arg = _get_parameters_arg(algorithm=algorithm) + if params_arg: + command_args += params_arg extra_args = {} if not algorithm.run_args.requires_root: @@ -226,12 +253,11 @@ def _sanity_check_output( BraTSContainerException: If not enough output files exist """ - inputs = list(data_path.iterdir()) + # some algorithms create extra files in the data folder, so we only check for files starting with "BraTS" + # (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(inputs) == len(outputs) and len(outputs) > 0: - return - if len(outputs) < len(inputs): logger.error(f"Docker container output: \n\r{container_output}") raise BraTSContainerException( @@ -259,7 +285,7 @@ def run_docker( # ensure image is present, if not pull it _ensure_image(image=algorithm.run_args.docker_image) - # Log the message + additional_files_path = _get_additional_files_path(algorithm) # ensure output folder exists @@ -269,19 +295,18 @@ def run_docker( data_path=data_path, additional_files_path=additional_files_path, output_path=output_path, + parameters_path=PARAMETERS_DIR, ) logger.debug(f"Volume mappings: {volume_mappings}") - command_args, extra_args = _build_args( - algorithm=algorithm, additional_files_path=additional_files_path - ) + command_args, extra_args = _build_args(algorithm=algorithm) logger.debug(f"Command args: {command_args}, Extra args: {extra_args}") # device setup device_requests = _handle_device_requests( algorithm=algorithm, cuda_devices=cuda_devices, force_cpu=force_cpu ) - logger.debug(f"Device requests: {device_requests}") + logger.debug(f"GPU Device requests: {device_requests}") # Run the container logger.info(f"{'Starting inference'}") diff --git a/brats/weights.py b/brats/weights.py index 1cdd7f2..4966f1f 100644 --- a/brats/weights.py +++ b/brats/weights.py @@ -11,8 +11,7 @@ from loguru import logger from tqdm import tqdm -ZENODO_RECORD_BASE_URL = "https://zenodo.org/api/records" -WEIGHTS_FOLDER = Path(__file__).parent / "weights" +from brats.constants import WEIGHTS_FOLDER, ZENODO_RECORD_BASE_URL def get_dummy_weights_path() -> Path: @@ -115,6 +114,7 @@ def _get_zenodo_metadata_and_archive_url(record_id: str) -> Dict | None: logger.error( f"Cant find model weights for record_id '{record_id}' on Zenodo. Exiting..." ) + # TODO add proper exit exception data = response.json() return data["metadata"], data["links"]["archive"] diff --git a/run.py b/run.py index d6bd9db..b62a661 100644 --- a/run.py +++ b/run.py @@ -1,21 +1,50 @@ from pathlib import Path -from brats import AdultGliomaSegmenter, MeningiomaSegmenter, PediatricSegmenter +from brats import ( + AdultGliomaSegmenter, + MeningiomaSegmenter, + PediatricSegmenter, + AfricaSegmenter, +) from brats.constants import ( AdultGliomaAlgorithms, MeningiomaAlgorithms, PediatricAlgorithms, + AfricaAlgorithms, ) +alg = AfricaAlgorithms.BraTS23_1 +segmenter = AfricaSegmenter(algorithm=alg, cuda_devices="0") +base = Path("/home/marcelrosier/brats_data/africa/BraTS-SSA-00126-000") +segmenter.infer_single( + t1c=base / "BraTS-SSA-00126-000-t1c.nii.gz", + t1n=base / "BraTS-SSA-00126-000-t1n.nii.gz", + t2f=base / "BraTS-SSA-00126-000-t2f.nii.gz", + t2w=base / "BraTS-SSA-00126-000-t2w.nii.gz", + output_file=f"africa_out/seg-{alg.value}.nii.gz", + log_file=f"africa_out/log-{alg.value}.log", +) + # alg = AdultGliomaAlgorithms.BraTS23_1 -# inferer = AdultGliomaInferer(algorithm=alg, cuda_devices="4") -# base = Path("/home/ivan_marcel/test_data/GLI/BraTS-GLI-00001-000/") +# inferer = AdultGliomaSegmenter(algorithm=alg, cuda_devices="0") + +# base = Path("/home/marcelrosier/brats_data/adult_glioma/BraTS-GLI-00001-000") # inferer.infer_single( -# t1c=base / "BraTS-GLI-00001-000-t1c.nii.gz", -# t1n=base / "BraTS-GLI-00001-000-t1n.nii.gz", -# t2f=base / "BraTS-GLI-00001-000-t2f.nii.gz", -# t2w=base / "BraTS-GLI-00001-000-t2w.nii.gz", +# t1c=base / "t1c.nii.gz", +# t1n=base / "t1n.nii.gz", +# t2f=base / "t2f.nii.gz", +# t2w=base / "t2w.nii.gz", # output_file=f"single_out/seg-{alg.value}.nii.gz", +# log_file=f"single_out/log-{alg.value}.log", +# ) +# base = Path("/home/marcelrosier/brats_data/oslo/16236606/day_0000") +# inferer.infer_single( +# t1c=base / "_t1c.nii.gz", +# t1n=base / "_t1.nii.gz", +# t2f=base / "_fla.nii.gz", +# t2w=base / "_t2.nii.gz", +# output_file=f"single_out/seg-{alg.value}.nii.gz", +# log_file=f"single_out/oslo-log-{alg.value}.log", # ) # import time @@ -35,10 +64,10 @@ # # pediatric -alg = PediatricAlgorithms.BraTS23_3 -segmenter = PediatricSegmenter(algorithm=alg, cuda_devices="4") +# alg = PediatricAlgorithms.BraTS23_3 +# segmenter = PediatricSegmenter(algorithm=alg, cuda_devices="4") -base = Path("/home/ivan_marcel/test_data/PED/BraTS-PED-00030-000") +# base = Path("/home/ivan_marcel/test_data/PED/BraTS-PED-00030-000") # segmenter.infer_single( # t1c=base / "BraTS-PED-00030-000-t1c.nii.gz", # t1n=base / "BraTS-PED-00030-000-t1n.nii.gz", @@ -47,8 +76,8 @@ # output_file=f"single_out/seg-{alg.value}.nii.gz", # log_file=f"single_out/log-{alg.value}.txt", # ) -segmenter.infer_batch( - data_folder=base.parent, - output_folder=Path("batch_out"), - log_file=Path("batch_out/log.log"), -) +# segmenter.infer_batch( +# data_folder=base.parent, +# output_folder=Path("batch_out"), +# log_file=Path("batch_out/log.log"), +# ) diff --git a/tests/test_algorithm_config.py b/tests/test_algorithm_config.py index 21d5b86..2a7b8b1 100644 --- a/tests/test_algorithm_config.py +++ b/tests/test_algorithm_config.py @@ -2,16 +2,14 @@ from pathlib import Path from brats.algorithm_config import load_algorithms - -PACKAGE_DIR = Path(__file__).parent.parent +from brats.constants import ALGORITHM_DIR class TestAlgorithmConfig(unittest.TestCase): def test_configs_valid(self): - algorithms_folder = PACKAGE_DIR / "brats" / "algorithms" - configs = algorithms_folder.iterdir() + configs = [f for f in ALGORITHM_DIR.iterdir() if f.is_file()] for config in configs: try: