Example: inference with pretrained nnU-Net models
This is a step-by-step example on 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 withnnUNet_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 outpput 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 donw 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.