BraTS 0.0.0.dev0__py3-none-any.whl

brats/core/docker.py

@@ -0,0 +1,474 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import docker
+import nibabel as nib
+import numpy as np
+from docker.errors import DockerException
+from loguru import logger
+from rich.console import Console
+from rich.progress import Progress
+from rich.table import Table
+from rich import box
+
+from brats.constants import DUMMY_PARAMETERS, PARAMETERS_DIR, PACKAGE_CITATION
+from brats.utils.algorithm_config import AlgorithmData
+from brats.utils.exceptions import (
+    AlgorithmNotCPUCompatibleException,
+    BraTSContainerException,
+)
+from brats.utils.zenodo import check_additional_files_path, get_dummy_path
+
+try:
+    client = docker.from_env()
+except DockerException as e:
+    logger.error(
+        f"Failed to connect to docker daemon. Please make sure docker is installed and running. Error: {e}"
+    )
+    # not aborting since this happens during read the docs builds. not a great solution tbf
+
+
+def _show_docker_pull_progress(tasks: Dict, progress: Progress, line: Dict):
+    """Show the progress of a docker pull operation.
+
+    Args:
+        tasks (Dict): The tasks to update
+        progress (Progress): The progress bar to update
+        line (Dict): the next line from docker.client.api.pull stream
+    """
+    status = line.get("status")
+    layer_id = line.get("id", "unknown")
+
+    if status in {"Downloading", "Extracting"}:
+        task_key = f"[{status} {layer_id}]"
+
+        progress_detail = line.get("progressDetail") or {}
+        total = progress_detail.get("total")
+        current = progress_detail.get("current")
+
+        # Case 1: we know total size -> normal progress bar
+        if total:
+            if task_key not in tasks:
+                tasks[task_key] = progress.add_task(f"{task_key}", total=total)
+            progress.update(tasks[task_key], completed=current or 0)
+
+        # Case 2: no total/current -> animated bar showing indefinite progress
+        else:
+            if task_key not in tasks:
+                tasks[task_key] = progress.add_task(
+                    f"{task_key}", total=None  # total=None means indeterminate
+                )
+            progress.update(tasks[task_key], advance=0.1)
+
+
+def _ensure_image(image: str):
+    """Ensure the docker image is present on the system. If not, pull it.
+
+    Args:
+        image (str): The docker image to pull
+    """
+    if not client.images.list(name=image):
+        logger.info(f"Pulling docker image {image}")
+        tasks = {}
+        with Progress() as progress:
+            resp = client.api.pull(image, stream=True, decode=True)
+            for line in resp:
+                _show_docker_pull_progress(tasks=tasks, progress=progress, line=line)
+
+
+def _is_cuda_available() -> bool:
+    """Check if CUDA is available on the system by trying to run nvidia-smi."""
+    try:
+        # Attempt to run `nvidia-smi` to check for CUDA.
+        # This command should run successfully if NVIDIA drivers are installed and GPUs are present.
+        subprocess.run(
+            ["nvidia-smi"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=True,
+        )
+        return True
+    except:
+        return False
+
+
+def _handle_device_requests(
+    algorithm: AlgorithmData, cuda_devices: str, force_cpu: bool
+) -> List[docker.types.DeviceRequest]:
+    """Handle the device requests for the docker container (request cuda or cpu).
+
+    Args:
+        algorithm (AlgorithmData): Algorithm data
+        cuda_devices (str): The CUDA devices to use
+        force_cpu (bool): Whether to force CPU execution
+    """
+    cuda_available = _is_cuda_available()
+    if not cuda_available or force_cpu:
+        if not algorithm.run_args.cpu_compatible:
+            cause = (
+                "User tried to force CPU execution but"
+                if cuda_available
+                else "No Cuda installation/ GPU was found and"
+            )
+            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 [
+        docker.types.DeviceRequest(device_ids=[cuda_devices], capabilities=[["gpu"]])
+    ]
+
+
+def _get_additional_files_path(algorithm: AlgorithmData) -> Path:
+    """Get the path to the additional files for @algorithm.
+
+    Args:
+        algorithm (AlgorithmData): The algorithm data
+
+    Returns:
+        Path to the additional files
+    """
+    # 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 additional_files are directly specified a dummy additional_files folder will be mounted
+        return get_dummy_path()
+
+
+def _get_volume_mappings_mlcube(
+    data_path: Path,
+    additional_files_path: Path,
+    output_path: Path,
+    parameters_path: Path,
+) -> Dict:
+    """Get the volume mappings for the docker container.
+
+    Args:
+        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, 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, parameters_path]
+        )
+    }
+
+
+def _get_volume_mappings_docker_only(
+    data_path: Path,
+    output_path: Path,
+) -> Dict[str, Dict[str, str]]:
+    """Get the volume mappings for the docker container.
+
+    Args:
+        data_path (Path): The path to the input data
+        output_path (Path): The path to save the output
+
+    Returns:
+        Dict: The volume mappings
+    """
+    return {
+        str(data_path.absolute()): {"bind": "/input", "mode": "rw"},
+        str(output_path.absolute()): {"bind": "/output", "mode": "rw"},
+    }
+
+
+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_command_args(
+    algorithm: AlgorithmData,
+) -> str:
+    """Build the command arguments for the docker container.
+
+    Args:
+        algorithm (AlgorithmData): The algorithm data
+
+    Returns:
+        command_args: The command arguments
+    """
+    # Build command that will be run in the docker container
+    command_args = f"--data_path=/mlcube_io0 --output_path=/mlcube_io2"
+    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)
+    if params_arg:
+        command_args += params_arg
+
+    return command_args
+
+
+def _get_container_user(
+    algorithm: AlgorithmData,
+) -> Optional[str]:
+    """Build extra arguments for the docker container.
+
+    Args:
+        algorithm (AlgorithmData): The algorithm data
+
+    Returns:
+        Optional[str]: The user to run the container as or None if root is required
+
+    """
+
+    if not algorithm.run_args.requires_root:
+        # run the container as the current user to ensure written files are always owned by the user
+        # also overall better security-wise
+        return f"{os.getuid()}:{os.getgid()}"
+    return None
+
+
+def _observe_docker_output(container: docker.models.containers.Container) -> str:
+    """Observe the output of a running docker container and display a spinner. On Errors log container output.
+
+    Args:
+        container (docker.models.containers.Container): The container to observe
+    """
+    # capture the output
+    container_output = container.attach(
+        stdout=True, stderr=True, stream=True, logs=True
+    )
+
+    # Display spinner while the container is running
+    with Console().status("Running inference..."):
+        # Wait for the container to finish
+        exit_code = container.wait()
+        container_output = "\n\r".join(
+            [line.decode("utf-8", errors="replace") for line in container_output]
+        )
+        # Check if the container exited with an error
+        if exit_code["StatusCode"] != 0:
+            logger.error(f">> {container_output}")
+            raise BraTSContainerException(
+                "Container finished with an error:\n"
+                f"{'-'*80}\n{container_output}\n {'-'*80}\n"
+                "For further debugging please specify a log file to capture the output or enable console logging (brats.utils.logging.enable())"
+            )
+
+    return container_output
+
+
+def _sanity_check_output(
+    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
+    """
+
+    # 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(outputs) < len(inputs):
+        logger.error(f"Docker container output: \n\r{container_output}")
+        raise BraTSContainerException(
+            f"Not enough output files were created by the algorithm. Expected: {len(inputs)} Got: {len(outputs)}. Please check the logging output of the docker container for more information."
+        )
+
+    for i, output in enumerate(outputs, start=1):
+        content = nib.load(output).get_fdata()
+        if np.count_nonzero(content) == 0:
+            name = "<unknown>"
+            if internal_external_name_map is not None:
+                name_key = [
+                    k for k in internal_external_name_map.keys() if k in output.name
+                ]
+                if name_key:
+                    name = internal_external_name_map[name_key[0]]
+
+            logger.warning(
+                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."""
+            )
+
+
+def _log_algorithm_info(algorithm: AlgorithmData):
+    """Log information about the algorithm and citation reminder.
+
+    Args:
+        algorithm (AlgorithmData): algorithm data
+    """
+
+    # intentional prints! (to always display them even when logging is disabled)
+    console = Console()
+    console.rule("[bold red]Citation Reminder[/bold red]")
+    console.print(
+        "Please support our development by citing the relevant manuscripts for the used algorithm:\n"
+    )
+    table = Table(
+        show_header=False,
+        show_lines=True,
+        show_edge=False,
+        box=box.ASCII,
+    )
+    table.add_column("type", justify="right", style="cyan")
+    table.add_column("url", justify="left", style="white")
+    table.add_row("BraTS Package", PACKAGE_CITATION)
+    table.add_row(
+        f"Challenge ({algorithm.meta.challenge} {algorithm.meta.year})",
+        algorithm.meta.challenge_manuscript,
+    )
+    table.add_row(f"Algorithm ({algorithm.meta.authors})", algorithm.meta.paper)
+    if algorithm.meta.dataset_manuscript:
+        table.add_row(f"Dataset", algorithm.meta.dataset_manuscript)
+
+    console.print(table)
+    console.rule()
+
+    logger.opt(colors=True).info(
+        f"Running algorithm: <light-green> BraTS {algorithm.meta.year} {algorithm.meta.challenge} [{algorithm.meta.rank} place]</>"
+    )
+    logger.debug(f"Docker image: {algorithm.run_args.docker_image}")
+
+
+def run_container(
+    algorithm: AlgorithmData,
+    data_path: Path,
+    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.
+
+    Args:
+        algorithm (AlgorithmData): The data of the algorithm to run
+        data_path (Path | str): The path to the input data
+        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)
+
+    # ensure image is present, if not pull it
+    _ensure_image(image=algorithm.run_args.docker_image)
+
+    additional_files_path = _get_additional_files_path(algorithm)
+
+    # ensure output folder exists
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # device setup
+    device_requests = _handle_device_requests(
+        algorithm=algorithm,
+        cuda_devices=cuda_devices,
+        force_cpu=force_cpu,
+    )
+    logger.debug(f"GPU Device requests: {device_requests}")
+
+    user = _get_container_user(algorithm=algorithm)
+    logger.debug(f"Container user: {user if user else 'root (required by algorithm)'}")
+
+    logger.info(f"{'Starting inference'}")
+    start_time = time.time()
+
+    if algorithm.meta.year <= 2024:
+        volume_mappings = _get_volume_mappings_mlcube(
+            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 = _build_command_args(algorithm=algorithm)
+        logger.debug(f"Command args: {command_args}")
+
+        # Run the container
+        container = client.containers.run(
+            image=algorithm.run_args.docker_image,
+            volumes=volume_mappings,
+            device_requests=device_requests,
+            command=f"infer {command_args}",
+            network_mode="none",
+            detach=True,
+            remove=True,
+            shm_size=algorithm.run_args.shm_size,
+            user=user,
+        )
+    else:
+        volume_mappings = _get_volume_mappings_docker_only(
+            data_path=data_path,
+            output_path=output_path,
+        )
+        logger.debug(f"Volume mappings: {volume_mappings}")
+        container = client.containers.run(
+            image=algorithm.run_args.docker_image,
+            volumes=volume_mappings,
+            device_requests=device_requests,
+            network_mode="none",
+            detach=True,
+            remove=True,
+            shm_size=algorithm.run_args.shm_size,
+            user=user,
+        )
+
+    container_output = _observe_docker_output(container=container)
+    _sanity_check_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}")
+
+    logger.info(f"Finished inference in {time.time() - start_time:.2f} seconds")

brats/core/inpainting_algorithms.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from pathlib import Path
+import shutil
+import sys
+from typing import Dict, Optional
+
+from loguru import logger
+
+from brats.core.brats_algorithm import BraTSAlgorithm
+from brats.constants import INPAINTING_ALGORITHMS, InpaintingAlgorithms, Task, Backends
+from brats.utils.data_handling import input_sanity_check
+
+
+class Inpainter(BraTSAlgorithm):
+
+    def __init__(
+        self,
+        algorithm: InpaintingAlgorithms = InpaintingAlgorithms.BraTS23_1,
+        cuda_devices: str | None = "0",
+        force_cpu: bool = False,
+    ):
+        super().__init__(
+            algorithm=algorithm,
+            algorithms_file_path=INPAINTING_ALGORITHMS,
+            task=Task.INPAINTING,
+            cuda_devices=cuda_devices,
+            force_cpu=force_cpu,
+        )
+
+    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.
+
+        Args:
+            data_folder (Path): Path to the data folder
+            subject_id (str): Subject ID
+            inputs (dict[str, Path | str]): Dictionary with the input data
+            subject_modality_separator (str): Separator between the subject ID and the modality
+        """
+
+        subject_folder = data_folder / subject_id
+        subject_folder.mkdir(parents=True, exist_ok=True)
+        # TODO: investigate usage of symlinks (might cause issues on windows and would probably require different volume handling)
+        t1n, mask = inputs["t1n"], inputs["mask"]
+        try:
+            shutil.copy(
+                t1n,
+                subject_folder
+                / f"{subject_id}{subject_modality_separator}t1n-voided.nii.gz",
+            )
+            shutil.copy(
+                mask,
+                subject_folder / f"{subject_id}{subject_modality_separator}mask.nii.gz",
+            )
+        except FileNotFoundError as e:
+            logger.error(f"Error while standardizing files: {e}")
+            raise
+
+        # sanity check inputs
+        input_sanity_check(t1n=t1n, mask=mask)
+
+    def _standardize_batch_inputs(
+        self, data_folder: Path, subjects: list[Path], input_name_schema: str
+    ) -> Dict[str, str]:
+        """Standardize the input images for a list of subjects to match requirements of all algorithms and save them in @tmp_data_folder/@subject_id.
+
+        Args:
+            subjects (List[Path]): List of subject folders, each with a voided t1n and a mask image
+            data_folder (Path): Parent folder where the subject folders will be created
+            input_name_schema (str): Schema to be used for the subject folder and filenames depending on the BraTS Challenge
+
+        Returns:
+            Dict[str, str]: Dictionary mapping internal name (in standardized format) to external subject name provided by user
+        """
+        internal_external_name_map = {}
+        for i, subject in enumerate(subjects):
+            internal_name = input_name_schema.format(id=i)
+            internal_external_name_map[internal_name] = subject.name
+            # TODO Add support for .nii files
+
+            self._standardize_single_inputs(
+                data_folder=data_folder,
+                subject_id=internal_name,
+                inputs={
+                    "t1n": subject / f"{subject.name}-t1n-voided.nii.gz",
+                    "mask": subject / f"{subject.name}-mask.nii.gz",
+                },
+                subject_modality_separator=self.algorithm.run_args.subject_modality_separator,
+            )
+        return internal_external_name_map
+
+    def infer_single(
+        self,
+        t1n: Path | str,
+        mask: Path | str,
+        output_file: Path | str,
+        log_file: Optional[Path | str] = None,
+        backend: Optional[Backends] = Backends.DOCKER,
+    ) -> None:
+        """Perform inpainting task on a single subject with the provided images and save the result to the output file.
+
+        Args:
+            t1n (Path | str): Path to the voided T1n image
+            mask (Path | str): Path to the mask image
+            output_file (Path | str): Path to save the segmentation
+            log_file (Path | str, optional): Save logs to this file
+            backend (Backends, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+        """
+
+        self._infer_single(
+            inputs={"t1n": t1n, "mask": mask},
+            output_file=output_file,
+            log_file=log_file,
+            backend=backend,
+        )
+
+    def infer_batch(
+        self,
+        data_folder: Path | str,
+        output_folder: Path | str,
+        log_file: Path | str | None = None,
+        backend: Optional[Backends] = Backends.DOCKER,
+    ) -> None:
+        """Perform inpainting on a batch of subjects with the provided images and save the results to the output folder. \n
+        Requires the following structure:\n
+        data_folder\n
+        ┣ A\n
+        ┃ ┣ A-t1n-voided.nii.gz\n
+        ┃ ┣ A-mask.nii.gz\n
+        ┣ B\n
+        ┃ ┣ B-t1n-voided.nii.gz\n
+        ┃ ┣ B-mask.nii.gz\n
+        ┣ C ...\n
+
+
+        Args:
+            data_folder (Path | str): Folder containing the subjects with required structure
+            output_folder (Path | str): Output folder to save the segmentations
+            log_file (Path | str, optional): Save logs to this file
+            backend (Backends, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+        """
+
+        return self._infer_batch(
+            data_folder=data_folder,
+            output_folder=output_folder,
+            log_file=log_file,
+            backend=backend,
+        )