inference.mednext.runner

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