inference.swinunetr.runner

 1"""
 2### SwinUNETR Inference Runner
 3
 4Provides functionalities related to running the SwinUNETR inference on the input data.
 5"""
 6
 7
 8import subprocess
 9from pathlib import Path
10import torch
11from typing import Any, List
12
13
14def run_infer_swinunetr(
15    input_path: Path,
16    output_folder: Path,
17    challenge: Any,
18    folds: List[int] = [0, 1, 2, 3, 4]
19) -> List[Path]:
20    """Run SwinUNETR inference on the provided input data.
21
22    This function executes the SwinUNETR inference script for each specified fold,
23    using the appropriate pretrained model weights, and collects the paths to the
24    resulting probability npz files.
25
26    Args:
27        input_path (Path): Path to the folder containing the 4 input files.
28        output_folder (Path): Path to the folder to store the npz files.
29        challenge (Any): The challenge identifier or related parameter.
30        folds (List[int], optional): List of fold indices to process. Defaults to [0, 1, 2, 3, 4].
31
32    Returns:
33        List[Path]: List of paths to the npz files generated by inference.
34    """
35    pretrained_name = 'best_model.pt'
36    npz_folder_list: List[Path] = []
37
38    for fold in folds:
39        # Determine the epoch based on fold
40        if fold in [3, 4]:
41            e = 1000
42        else:
43            e = 650
44
45        # Define the path to the pretrained model
46        pretrained_path = Path(f"/tmp/swinunetr_e{e}_f{fold}_b1p4/")
47
48        # Load arguments from the pretrained model
49        args = torch.load(pretrained_path / pretrained_name)['args']
50
51        # Define the output directory for this fold
52        output_dir = output_folder / f"swinunetr3d_f{fold}"
53        npz_folder_list.append(output_dir)
54
55        # Construct the command to run the SwinUNETR inference
56        cmd = 'python swinunetr/inference.py'
57        cmd = ' '.join((cmd, f"--datadir='{str(input_path)}'"))
58        cmd = ' '.join((cmd, f"--exp_path='{str(output_dir)}'"))
59        cmd = ' '.join((cmd, f'--roi_x={args.roi_x}'))
60        cmd = ' '.join((cmd, f'--roi_y={args.roi_y}'))
61        cmd = ' '.join((cmd, f'--roi_z={args.roi_z}'))
62        cmd = ' '.join((cmd, f'--in_channels={args.in_channels}'))
63        cmd = ' '.join((cmd, f'--out_channels={args.out_channels}'))
64        cmd = ' '.join((cmd, '--spatial_dims=3'))
65        cmd = ' '.join((cmd, '--use_checkpoint'))
66        cmd = ' '.join((cmd, '--feature_size=48'))
67        cmd = ' '.join((cmd, '--infer_overlap=0.625'))
68        cmd = ' '.join((cmd, '--cacherate=1.0'))
69        cmd = ' '.join((cmd, '--workers=0'))
70        cmd = ' '.join((cmd, f'--pretrained_model_name={pretrained_name}'))
71        cmd = ' '.join((cmd, f'--pretrained_dir={str(pretrained_path)}'))
72        cmd = ' '.join((cmd, '--pred_label'))
73
74        print(cmd)
75        subprocess.run(cmd, shell=True)  # Executes the command in the shell
76
77    # Generate list of npz paths based on output folders
78    npz_path_list = [f / f"{input_path.name}-t1n.npz" for f in npz_folder_list]
79    return npz_path_list