Example: inference with pretrained nnU-Net models
This is a step-by-step example of how to run inference with pretrained nnU-Net models on the Prostate dataset of the Medical Segemtnation Decathlon.
Install nnU-Net by following the instructions here. Make sure to set all relevant paths, also see here. This step is necessary so that nnU-Net knows where to store trained models.
Download the Prostate dataset of the Medical Segmentation Decathlon from here. Then extract the archive to a destination of your choice.
We selected the Prostate dataset for this example because we have a utility script that converts the test data into the correct format.
Decathlon data come as 4D niftis. This is not compatible with nnU-Net (see dataset format specified here). Convert the Prostate dataset into the correct format with
nnUNet_convert_decathlon_task -i /xxx/Task05_Prostate
Note that
Task05_Prostate
must be the folder that has the three 'imagesTr', 'labelsTr', 'imagesTs' subfolders! The converted dataset can be found in$nnUNet_raw_data_base/nnUNet_raw_data
($nnUNet_raw_data_base is the folder for raw data that you specified during installation)Download the pretrained model using this command:
nnUNet_download_pretrained_model Task005_Prostate
The prostate dataset requires two image modalities as input. This is very much liKE RGB images have three color channels. nnU-Net recognizes modalities by the file ending: a single test case of the prostate dataset therefore consists of two files
case_0000.nii.gz
andcase_0001.nii.gz
. Each of these files is a 3D image. The file ending with 0000.nii.gz must always contain the T2 image and 0001.nii.gz the ADC image. Whenever you are using pretrained models, you can usennUNet_print_pretrained_model_info Task005_Prostate
to obtain information on which modality needs to get which number. The output for Prostate is the following:
Prostate Segmentation. Segmentation targets are peripheral and central zone, input modalities are 0: T2, 1: ADC. Also see Medical Segmentation Decathlon, http://medicaldecathlon.com/
The script we ran in 3) automatically converted the test data for us and stored them in
$nnUNet_raw_data_base/nnUNet_raw_data/Task005_Prostate/imagesTs
. Note that you need to do this conversion youself when using other than Medcial Segmentation Decathlon datasets. No worries. Doing this is easy (often as simple as appending a _0000 to the file name if only one input modality is required). Instructions can be found here here.You can now predict the Prostate test cases with the pretrained model. We exemplarily use the 3D full resoltion U-Net here:
nnUNet_predict -i $nnUNet_raw_data_base/nnUNet_raw_data/Task005_Prostate/imagesTs/ -o OUTPUT_DIRECTORY -t 5 -m 3d_fullres
Note that
-t 5
specifies the task with id 5 (which corresponds to the Prostate dataset). You can also give the full task nameTask005_Prostate
.OUTPUT_DIRECTORY
is where the resulting segmentations are saved.Predictions should be quite fast and you should be done within a couple of minutes. If you would like to speed it up (at the expense of a slightly lower segmentation quality) you can disable test time data augmentation by setting the
--disable_tta
flag (8x speedup). If this is still too slow for you, you can consider using only a single model instead of the ensemble by specifying-f 0
. This will use only the model trained on fold 0 of the cross-validation for another 5x speedup.If you want to use an ensemble of different U-Net configurations for inference, you need to run the following commands:
Prediction with 3d full resolution U-Net (this command is a little different than the one above).
nnUNet_predict -i $nnUNet_raw_data_base/nnUNet_raw_data/Task005_Prostate/imagesTs/ -o OUTPUT_DIRECTORY_3D -t 5 --save_npz -m 3d_fullres
Prediction with 2D U-Net
nnUNet_predict -i $nnUNet_raw_data_base/nnUNet_raw_data/Task005_Prostate/imagesTs/ -o OUTPUT_DIRECTORY_2D -t 5 --save_npz -m 2d
--save_npz
will tell nnU-Net to also store the softmax probabilities for ensembling.You can then merge the predictions with
nnUNet_ensemble -f OUTPUT_DIRECTORY_3D OUTPUT_DIRECTORY_2D -o OUTPUT_FOLDER_ENSEMBLE -pp POSTPROCESSING_FILE
This will merge the predictions from
OUTPUT_DIRECTORY_2D
andOUTPUT_DIRECTORY_3D
.-pp POSTPROCESSING_FILE
(optional!) is a file that gives nnU-Net information on how to postprocess the ensemble. These files were also downloaded as part of the pretrained model weights and are located atRESULTS_FOLDER/nnUNet/ensembles/ Task005_Prostate/ensemble_2d__nnUNetTrainerV2__nnUNetPlansv2.1--3d_fullres__nnUNetTrainerV2__nnUNetPlansv2.1/postprocessing.json
. We will make the postprocessing files more accessible in a future (soon!) release.