inference.nnunet.runner

  1"""
  2### nnUNet Inference Runner
  3
  4Provides functionalities related to running the nnUNet inference on the input data.
  5"""
  6
  7
  8import subprocess
  9from tqdm import tqdm
 10import os
 11from pathlib import Path
 12import argparse
 13from typing import List, Tuple, Dict, Optional
 14
 15
 16def maybe_make_dir(path: str) -> str:
 17    """
 18    Creates a directory at the specified path if it does not exist.
 19
 20    Args:
 21        path (str): Path to the directory to be created.
 22
 23    Returns:
 24        str: The path to the created or existing directory.
 25    """
 26    os.makedirs(path, exist_ok=True)
 27    return path
 28
 29
 30def set_env_paths(input_path: str, path: str) -> str:
 31    """
 32    Sets environment paths required for nnUNet predictions.
 33
 34    Args:
 35        input_path (str): Path to the input data.
 36        path (str): Path to the output directory.
 37
 38    Returns:
 39        str: A string containing environment variables for nnUNet.
 40    """
 41    # Define the path for preprocessed data
 42    preprocessed = maybe_make_dir(str(Path('/tmp/') / 'nnUNet_preprocessed'))
 43    
 44    # Return environment variable settings as a single string
 45    return f"nnUNet_raw='{input_path}' nnUNet_preprocessed='{preprocessed}' nnUNet_results='/tmp/'"
 46
 47
 48def get_dataset_info(challenge_name: str) -> Tuple[str, str, str, str]:
 49    """
 50    Returns the corresponding dataset name, trainer, plans, and configuration based on the provided challenge name.
 51
 52    Args:
 53        challenge_name (str): Name of the challenge (e.g., BraTS-PED, BraTS-MET).
 54
 55    Returns:
 56        Tuple[str, str, str, str]: 
 57            - dataset_name: Name of the dataset.
 58            - trainer: Trainer configuration.
 59            - plans: Plans configuration.
 60            - configuration: nnUNet configuration.
 61
 62    Raises:
 63        Exception: If the challenge name is not compatible.
 64    """
 65    if challenge_name == "BraTS-PED":
 66        dataset_name = "Dataset021_BraTS2024-PED"
 67        trainer = "nnUNetTrainer_200epochs"
 68        plans = "nnUNetPlans"
 69        configuration = "3d_fullres"
 70    elif challenge_name == "BraTS-SSA":
 71        dataset_name = "Dataset022_BraTS2024-SSA"
 72        trainer = "nnUNetTrainer_150epochs_CustomSplit_stratified_pretrained_GLI"
 73        plans = "nnUNetPlans_pretrained_GLI"
 74        configuration = "3d_fullres"
 75    elif challenge_name == "BraTS-GLI":
 76        dataset_name = "Dataset023_BraTS2024-GLI"
 77        trainer = "nnUNetTrainer_200epochs"
 78        plans = "nnUNetPlans"
 79        configuration = "3d_fullres"
 80    elif challenge_name == "BraTS-MEN-RT":
 81        dataset_name = "Dataset024_BraTS2024-MEN-RT"
 82        trainer = "nnUNetTrainer_200epochs"
 83        plans = "nnUNetPlans"
 84        configuration = "3d_fullres"
 85    elif challenge_name == "BraTS-MET":
 86        dataset_name = "Dataset026_BraTS2024-MET"
 87        trainer = "nnUNetTrainer_200epochs"
 88        plans = "nnUNetPlans"
 89        configuration = "3d_fullres"
 90    else:
 91        raise Exception("Challenge name not compatible.")
 92    
 93    return dataset_name, trainer, plans, configuration
 94
 95
 96def run_infer_nnunet(
 97    input_folder: str,
 98    output_folder: str,
 99    challenge_name: str,
100    name: str,
101    folds: List[int] = [0, 1, 2, 3, 4],
102    save_npz: bool = True,
103    ensemble: bool = True
104) -> List[Path]:
105    """
106    Runs nnUNet inference on the provided input data.
107
108    This function executes the nnUNet inference command either as an ensemble
109    across all folds or individually per specified fold. It extracts the
110    predictions and saves them as .npz files.
111
112    Args:
113        input_folder (str): Path to the input folder containing prediction files.
114        output_folder (str): Path to the folder where output predictions will be stored.
115        challenge_name (str): Name of the challenge (e.g., BraTS-PED, BraTS-MET).
116        name (str): Base name for the output prediction files.
117        folds (List[int], optional): List of fold indices to process. Defaults to [0, 1, 2, 3, 4].
118        save_npz (bool, optional): Whether to save the prediction probabilities as .npz files. Defaults to True.
119        ensemble (bool, optional): Whether to perform ensemble prediction across all folds. Defaults to True.
120
121    Returns:
122        List[Path]: A list of paths to the saved .npz files from each fold.
123    """
124    # Retrieve dataset information based on the challenge name
125    dataset_name, trainer, plans, configuration = get_dataset_info(challenge_name)
126    
127    # Set environment paths required for nnUNet
128    env_set = set_env_paths(input_folder, output_folder)
129
130    if ensemble:
131        # Define the output directory for ensemble predictions
132        output_folder_fold = os.path.join(output_folder, "ens")
133        print(f"Running nnUNet inference with all folds (ensemble)..")
134        
135        # Construct the nnUNet inference command for ensemble
136        cmd = (
137            f"{env_set} nnUNetv2_predict "
138            f"-i '{input_folder}' "
139            f"-o '{output_folder_fold}' "
140            f"-d '{dataset_name}' "
141            f"-c '{configuration}' "
142            f"-tr '{trainer}' "
143            f"-p '{plans}'"
144        )
145        
146        if save_npz:
147            cmd += " --save_probabilities"
148        
149        # Execute the inference command
150        subprocess.run(cmd, shell=True)
151        
152        # Return the path to the ensemble .npz file
153        return [Path(os.path.join(output_folder_fold, f"{name}.npz"))]
154    else:
155        # Initialize a list to store paths to .npz files for each fold
156        npz_path_list: List[Path] = [] 
157        
158        # Iterate over each specified fold and perform inference
159        for fold in tqdm(folds, desc="Processing Folds"):
160            output_folder_fold = os.path.join(output_folder, f"fold_{fold}")
161            print(f"Running nnUNet inference for fold {fold}")
162            
163            # Construct the nnUNet inference command for the current fold
164            cmd = (
165                f"{env_set} nnUNetv2_predict "
166                f"-i '{input_folder}' "
167                f"-o '{output_folder_fold}' "
168                f"-d '{dataset_name}' "
169                f"-c '{configuration}' "
170                f"-tr '{trainer}' "
171                f"-p '{plans}' "
172                f"-f '{fold}'"
173            )
174            
175            if save_npz:
176                cmd += " --save_probabilities"
177            
178            # Execute the inference command for the current fold
179            subprocess.run(cmd, shell=True)  # Executes the command in the shell
180            
181            # Append the path to the saved .npz file
182            npz_path_list.append(Path(os.path.join(output_folder_fold, f"{name}.npz")))
183    
184        return npz_path_list