inference.runner_ped

  1"""
  2### Pediatric Brain Tumor Segmentation Pipeline
  3
  4Provides functionalities related to running the pipeline for pediatric brain tumor segmentation in MRI.
  5"""
  6
  7
  8from pathlib import Path
  9import os
 10import shutil
 11import tempfile
 12from typing import Tuple, List, Dict, Any
 13import pandas as pd
 14
 15from .nnunet.install_model import install_model_from_zip
 16from .ensembler.ped_weighted_ensemble import ped_ensemble
 17from .nnunet.runner import run_infer_nnunet
 18from .mednext.runner import run_infer_mednext
 19from .swinunetr.runner import run_infer_swinunetr
 20from .pp_cluster.infer import get_cluster, get_cluster_artifacts
 21from .radiomics.feature_extraction_v2 import extract_all, save_json, load_json
 22from .postproc.postprocess_cc import remove_small_component
 23from .postproc.postprocess_lblredef import label_redefinition
 24
 25# Task identifier
 26TASK: str = 'BraTS-PED'
 27
 28# Read the weights_cv.json
 29# Determine the directory where this script is located
 30script_dir = Path(__file__).resolve().parent
 31    
 32json_path: Path = script_dir / "ensembler" / "weights_cv.json"
 33ENSEMBLE_WEIGHTS: Dict[str, float] = load_json(json_path)[TASK]
 34NAME_MAPPER: Dict[str, str] = load_json(script_dir / "weights" / "name.json")
 35CLUSTER_ARTIFACT: Any = get_cluster_artifacts(TASK)
 36THRESHOLD_FILE_CC: Path = script_dir / "postproc" / "cc_thresholds_cv.json"
 37THRESHOLD_FILE_LBLDEF: Path = script_dir / "postproc" / "lblredef_thresholds_cv.json"
 38
 39# Configuration constants
 40CONSTANTS: Dict[str, str] = {
 41    'pyradiomics_paramfile': str(script_dir / "radiomics" / "params.yaml"),
 42    'nnunet_model_path': str(script_dir / "weights" / "BraTS2024_PEDs_nnunetv2_model.zip"),
 43    'mednext_model_path': str(script_dir / "weights" / "BraTS2024_PEDs_MedNeXt_model.zip"),
 44    'swinunetr_model_path': str(script_dir / "weights" / "swinunetr_peds_trunc.zip"),
 45}
 46
 47
 48def maybe_make_dir(path: Path) -> Path:
 49    """Create a directory if it does not exist.
 50
 51    Args:
 52        path (Path): The path of the directory to create.
 53
 54    Returns:
 55        Path: The path to the created directory.
 56    """
 57    os.makedirs(path, exist_ok=True)
 58    return Path(path)
 59
 60
 61def lbl_redefination(
 62    task: str,
 63    threshold_file: Path,
 64    input_dir: Path,
 65    df_radiomics: pd.DataFrame,
 66    out_dir: Path
 67) -> None:
 68    """Perform label redefinition by copying files from input to output directory.
 69
 70    Args:
 71        task (str): The task identifier.
 72        threshold_file (Path): Path to the threshold file for label redefinition.
 73        input_dir (Path): Directory containing input files.
 74        df_radiomics (pd.DataFrame): DataFrame containing radiomics data.
 75        out_dir (Path): Directory to store output files.
 76    """
 77    # Copy files from input_dir to out_dir
 78    for file in input_dir.iterdir():
 79        shutil.copy(file, out_dir)
 80
 81
 82def postprocess_single(
 83    input_dir: Path,
 84    seg_dir: Path,
 85    out_dir: Path
 86) -> Path:
 87    """Postprocess a single segmentation case.
 88
 89    This function performs radiomics extraction, cluster determination, removes small components,
 90    and redefines labels based on the processed data.
 91
 92    Args:
 93        input_dir (Path): Directory containing the input data.
 94        seg_dir (Path): Directory containing the segmentation results.
 95        out_dir (Path): Directory to store the postprocessed segmentation.
 96
 97    Returns:
 98        Path: Path to the postprocessed segmentation file.
 99    """
100    case_name: str = input_dir.name
101    seg_path: Path = seg_dir / f"{case_name}.nii.gz"
102    out_path: Path = maybe_make_dir(out_dir) / f"{case_name}.nii.gz"
103
104    # Compute radiomics features
105    try:
106        df_radiomics: pd.DataFrame = extract_all(
107            CONSTANTS['pyradiomics_paramfile'],
108            input_dir.parent,
109            case_name,
110            seg_path,
111            1,
112            region='wt',
113            tmpp='/tmp/',
114            seg_suffix='',
115            sequences=['-t1n', '-t1c', '-t2w', '-t2f']
116        )
117        # Determine the cluster based on radiomics features
118        cluster: int = get_cluster(df_radiomics, CLUSTER_ARTIFACT)[0]
119    except Exception as e:
120        print(f"Error in radiomics extraction: {e}")
121        df_radiomics = pd.DataFrame([{"StudyID": case_name}])
122        cluster = 4
123
124    df_radiomics['cluster'] = int(cluster)
125
126    # Remove disconnected regions from segmentation
127    removed_cc_dir: Path = maybe_make_dir(seg_dir / 'remove_cc')
128    remove_small_component(
129        TASK,
130        THRESHOLD_FILE_CC,
131        seg_dir,
132        df_radiomics,
133        removed_cc_dir
134    )
135
136    # Redefine labels based on the postprocessed segmentation
137    label_redefinition(
138        TASK,
139        THRESHOLD_FILE_LBLDEF,
140        removed_cc_dir,
141        df_radiomics,
142        out_dir
143    )
144
145    return out_path
146
147
148def infer_single(
149    input_path: Path,
150    out_dir: Path
151) -> None:
152    """Perform inference on a single input directory.
153
154    This function sets up a temporary directory, copies input files, runs inference models,
155    ensembles the predictions, and postprocesses the segmentation.
156
157    Args:
158        input_path (Path): Input directory containing the 4 nii.gz files.
159        out_dir (Path): Output directory to store the segmentation results.
160    """
161    with tempfile.TemporaryDirectory() as tmpdirname:
162        temp_dir: Path = Path(tmpdirname)
163        print(f'Storing artifacts in temporary directory {temp_dir}')
164
165        input_folder_raw: Path = maybe_make_dir(temp_dir / 'inp')
166        name: str = input_path.name
167        print(f'Processing case: {name}')
168
169        shutil.copytree(input_path, input_folder_raw / name)
170
171        # Rename files based on NAME_MAPPER
172        for key, val in NAME_MAPPER.items():
173            src: Path = input_folder_raw / name / f'{name}{key}'
174            dest: Path = input_folder_raw / name / f'{name}{val}'
175            os.rename(src, dest)
176            one_image: Path = dest
177
178        # Run inference using nnUNet
179        nnunet_npz_path_list: List[Path] = run_infer_nnunet(
180            input_folder_raw / name,
181            maybe_make_dir(temp_dir / 'nnunet'),
182            TASK,
183            name,
184            ensemble=False
185        )
186
187        # Run inference using MedNeXt
188        mednext_npz_path_list: List[Path]
189        mednext_pkl_path_list: List[Path]
190        mednext_npz_path_list, mednext_pkl_path_list = run_infer_mednext(
191            input_folder_raw / name,
192            maybe_make_dir(temp_dir / 'mednext'),
193            TASK,
194            name,
195            ensemble=False
196        )
197
198        # Run inference using SwinUNETR
199        swinunetr_npz_path_list: List[Path] = run_infer_swinunetr(
200            Path(input_path),
201            maybe_make_dir(temp_dir / 'swinunetr'),
202            TASK
203        )
204
205        # Ensemble the predictions
206        ensemble_folder: Path = maybe_make_dir(input_folder_raw / 'ensemble')
207        ensembled_pred_nii_path: Path = ped_ensemble(
208            swinunetr_npz_path_list,
209            nnunet_npz_path_list,
210            mednext_npz_path_list,
211            mednext_pkl_path_list,
212            ensemble_folder,
213            one_image,
214            weights=[
215                ENSEMBLE_WEIGHTS['SwinUNETR'],
216                ENSEMBLE_WEIGHTS['nnUNetv2'],
217                ENSEMBLE_WEIGHTS['MedNeXt']
218            ]
219        )
220
221        # Postprocess the ensembled segmentation
222        radiomics: Path = postprocess_single(
223            input_path,
224            ensemble_folder,
225            out_dir
226        )
227
228
229def batch_processor(
230    input_folder: str,
231    output_folder: str
232) -> None:
233    """Process a batch of input directories for inference.
234
235    Args:
236        input_folder (str): Path to the folder containing input directories.
237        output_folder (str): Path to the folder to store output segmentations.
238    """
239    for input_path in Path(input_folder).iterdir():
240        infer_single(input_path, Path(output_folder))
241
242
243def setup_model_weights() -> None:
244    """Install model weights from zip files."""
245    install_model_from_zip(CONSTANTS['nnunet_model_path'])
246    install_model_from_zip(CONSTANTS['swinunetr_model_path'])
247    install_model_from_zip(CONSTANTS['mednext_model_path'], mednext=True)
248
249
250def batch_postprocess(
251    input_folder: str,
252    seg_folder: str,
253    output_folder: str
254) -> int:
255    """Postprocess a batch of segmentation results.
256
257    Args:
258        input_folder (str): Path to the folder containing input directories.
259        seg_folder (str): Path to the folder containing segmentation results.
260        output_folder (str): Path to the folder to store postprocessed results.
261
262    Returns:
263        int: Number of processed cases.
264    """
265    l_radiomics: List[Path] = []
266    l_path: List[Path] = [x for x in Path(input_folder).iterdir() if x.is_dir()]
267    
268    for input_path in l_path:
269        radiomics: Path = postprocess_single(input_path, Path(seg_folder), Path(output_folder))
270        l_radiomics.append(radiomics)
271    
272    save_json(output_folder / 'val_radiomics.json', l_radiomics)
273    return len(list(Path(input_folder).iterdir()))
274
275
276if __name__ == "__main__":
277    setup_model_weights()
278    batch_processor('./ins', './outs')