nnU-Net : no-new-UNet (2024)

.ipynb format alternative for smooth execution of the codes.

Prateek Gupta, Indian Institute of Information Technology, Pune (prateekgupta16@alumni.iiitp.ac.in)

Kumar T. Rajamani, Institute of Medical Informatics, University of Lübeck, Germany (kumar.rajamani@uni-luebeck.de)

Mattias P. Heinrich, Institute of Medical Informatics, University of Lübeck, Germany (heinrich@imi.uni-luebeck.de)

nnU-Net is an open-source tool that can effectively be used out-of-the-box, rendering state of the art segmentation and catalyzing scientific progress as a framework for automated method design. It provides an end-to-end automated pipeline, which can be trained and inferred on any medical dataset for segmentation.

nnU-Net has outperformed state-of-the-art architecture(s) in the Medical Decathlon Challenge, comprising of 10 different datasets using an ensemble of the same U-Net architecture with an automated pipeline comprising of pre-processing, data augmentation and post-processing. It has set a new benchmark in the field of medical image segmentation, without the need to fine tuning a new architecture for every dataset individually. The pipeline itself takes care of the hyper-parameter tuning and requires no change in the network architecture to achieve start-of-the-art results.

• Working through nnU-net has it’s own hassles — manually renaming the files, putting the files in the right location, saving the data in the correct format. Each of these steps are time consuming and susceptible to human errors. Hence, this tutorial is written in a manner that will allow you to avoid the common errors.
• The paper detailing nnU-net has a long supplementary material for reference (55 pages research paper). We want to save your worthy time by summarizing the internal working of the nnU-net architecture to provide better understanding for first time users (having series of automation steps, it could take quite sometime to understand the internal workings.)
• We have tried to provide a single notebook to run through all the necessary steps as part of getting started with nnU-net. Though nnU-net has 3 different documentations detailing all the steps, our tutorial gives brief description of the parts relevant to users who just want to get started.

Though, nnU-Net takes care of most of the steps involved itself, sometimes getting it up and running for a custom dataset which is not in format mentioned here is difficult. Training the system on a colab notebook is again a challenge, as it requires a lot of manual steps. This tutorial will run you through all the necessary steps required to train your first nnU-net model along with ways to ensure the correctness of the procedure.

• Introduces the nnU-Net framework
• One-stop tutorial to train and test with different datasets
• Colab Notebook Based Tutorial (will work for local machines too).
• Folder creation, command execution from particular folder location done by the code.
• How to ingest the SCGM Challenge Dataset to the nnUNet pipeline.
• Data Visualization for Train, Test and Prediction

The nnU-Net pipeline uses heuristic rule to determine the data-dependent hyper-paramters, known as the “data fingerprint”, to ingest the training data. The blueprint parameters ( loss function, optimizer,architecture) and inferred parameters ( image resampling, normalization, batch and patch size) along with the data fingerprint generate pipeline fingerprints. Pipeline fingerprints produce network training for 2D, 3D and 3D-Cascade U-Net using the hyper-parameters determined so far. The ensemble of different network configuration(s), along with post-processing determines the best average Dice coefficient for the training data. The best configuration will then be used to produce the predictions for the test data. Details about individual component are described below.

It contains a set of heuristic rules to infer data-dependent hyper-parameters of the pipeline.

• Image size (i.e. number of voxels per spatial dimension) before and after cropping image
• Image Spacing (i.e. physical size of the voxels)
• Modalities (from metadata like CT, MRI)
• Number of classes for all images and total number of training cases

a. Architecture Template

• nnU-Net architecture closely follows the original U-Net and has recently proposed variations such as residual connection, attention mechanisms, squeeze and excitation, and dilated convolutions.
• Prefers large patch size rather than batch size.

b. Training Schedule

• 100 epochs; one epoch defined as iteration over 250 mini-batches.
• Stochastic gradient descent with Nestrov momentum.
• Loss is the combination of cross-entropy and dice loss.
• Oversampling to handle class imbalance.
• Data Augmentation : rotation, scaling, guassian noise, guassian blur, brightness, contrast, simulation of low resolution and gamma resolution.

c. Inference

• Images are predicted with a sliding window approach, where the window size equals the patch size used during training. Adjacent predictions overlap by half the size of a patch.
• To suppress stitching artifacts and reduce the influence of positions close to the borders, a Gaussian importance weighing is applied, increasing the weight of the center voxels in the softmax aggregation.

a. Intensity Normalization :

• Z-scoring for all modalities except CT.
• For CT, it follows a global normalization scheme, which uses 0.5 and 99.5 percentiles of the foreground voxels for clipping.

b. Resampling :

• To cope up with the heterogeneity in medical domain,
• Resamples all images to same target spacing using either third order spline, linear or nearest neighbor interpolation.

c. Adaptation of network topology, patch size and batch size : The network topology for all U-Net configurations is chosen on the basis of the median image size after resampling as well as the target spacing the images were resampled to.

d. Initialization : The patch size is initialized as the median image shape after resampling.

e. Architecture Topology:

• The architecture is configured by determining the number of down-sampling operations, performed until the feature map is reduced to 4 voxels or feature map space become anisotropic.
• High resolution axes are down-sampled separately until their resolution is within factor 2 of the lower resolution axis.
• Each axis is down-sampled individually, until the feature map constraints are triggered.
• The default kernel size for convolutions is 3×3×3 and 3×3 for 3D U-Net and 2D U-Net, respectively.

a. Ensembling and selection of U-Net configuration(s) :

• Automatically ensembles based on average foreground Dice coefficient computed via Cross-validation on training data to use for inference.
• Configuration model(s): Single models (2D, 3D fullres, 3D lowres or fullres U-Net of the cascade) or an ensemble of any two of these configurations.
• Models are ensembled by averaging softmax probabilities.

b. Postprocessing: Connected component-based postprocessing is used. All foreground classes are treated as one component, to improve the average foreground Dice coefficient and if it does not reduce the Dice coefficient for any of the classes, then nnU-Net builds on the outcome of this step and decides whether the same procedure should be performed for individual classes.

• 2D U-Net
• 3D U-Net Full Resolution
• 3D U-Net Cascaded : The first U-Net operates on downsampled images and the second is trained to refine the segmentation maps created by the former at full resolution.

Figure : for illustratation purposes. network architecture for ACDC dataset, from Medical Decathlon taken from Paper 1 mentioned below. Filters and no. of channels will remain same for any dataset.

SCGM Challenge, has a collection of healthy spinal cord images collected at 4 different sites. As part of the challenge, you need to automatically or semi-automatically segment the anatomical MR Images into background class (as 0) and grey matter class (as 1). Algorithms will be evaluated against manual segmentations from four trained raters (one from each site) in terms of segmentation accuracy and precision. Predicted images must be saved in NIFIT format, with same space and resolution as the provided data.

Note: Training images are labelled for Background class as (0), Grey Matter class as (1) and White Matter class as (2).

Train a nnUNet model for the Spinal Cord Grey Matter Segmentation Dataset and infer results . Same approach can be applied on other datasets, with dataset specific modifications.

• Setup (Always,CPU, GPU)
• Cloning the Repo’s (One-time, CPU, GPU)
• Install and import libraries (Everytime, CPU, GPU)
• Structure Dataset Folder (Everytime, CPU, GPU)
• Environment Variables (Everytime, CPU, GPU)
• Unzip, Rename, Train and Test Set (One-time, CPU, GPU)
• Dataset verification (One-time, Optional, CPU, GPU)
• Dataset Visualization (CPU,GPU)
• Training Code (GPU only)
• Inference Code (GPU only)
• Prediction Visualization (CPU,GPU)
• Submit Code for SCGM (Optional, CPU, GPU)

Important: Preferably use the .ipynb file to run the codes.

Note: For training and inference, you have to run through steps 1,3,4 everytime for colab users. For local machines, installing libraries and setting up environment variables again is not necessary.

Papers:

1) nnU-Net: Self-adapting Framework for U-Net-Based Med
ical Image Segmentation (arxiv)

2) Automated Design of Deep Learning Methods for Biomedical Image Segmentation (arxiv)

Github repository:

1) Original Repository — regularly maintained and changed by the authors.

2) Forked Repository — to ensure a version of nnUNet with which this tutorial works (as the original repository is constantly modified.)

(Colab users — Preferably use GPU runtime, but you can change to GPU runtime afterwards, as and when required)

(Colab Users: You must restart your runtime after installing the libraries.)

Custom Task Id starts at 101, to ensure that there will be no conflicts with downloaded pretrained models.

Colab Users: Everytime you re-run or restart your kernel always run until this point.

(Manual Task) Get the train and test data in zip form and place it in the folder: /nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/task_folder_name .

Apply to get Spinal Cord Grey Matter Challenge Dataset.

Code Below will take care of the following:

• Unzip train and test files
• Rename the train images and labels to match and placing in respective directories.
• Put testing data in the folder and remove text files
• Add modality at the end of each file, as nnU-net can train on multiple modalities together
• Create dataset.json

Note: For a new dataset, you may need to do few changes for training. Also, nnUNet works with .nii.gz files only.

We have 4 annotation of the same image, by different experts in the SCGM Challenge. ( Image , Ann1 ) and ( Image , Ann2 ) can be considered as a different image and label pairs. Hence, 4 copies of the training .nii.gz file is created with its mapping to the respective label name.

Verification

Note: In labelled image, yellow color represents white Matter and green-ish color represents grey matter.

nnU-Net stores a checkpoint every 50 epochs. If you need to continue a previous training, just add a -c to the training command.

Generic Training Commands:

nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD (additional options)

For 2D: nnUNet_train 2d nnUNetTrainerV2 TaskXXX_MYTASK FOLD

For 3D Full resolution: nnUNet_train 3d_fullres nnUNetTrainerV2 TaskXXX_MYTASK FOLD

For Cascaded 3D:

First Run lowres: nnUNet_train 3d_lowres nnUNetTrainerV2 TaskXXX_MYTASK FOLD

Then Run fullres: nnUNet_train 3d_cascade_fullres nnUNetTrainerV2CascadeFullRes TaskXXX_MYTASK FOLD

Training for 3D fullres with Trainer V2 and for Fold 0.

nnUNet_find_best_configuration will print inference commands you need to use. The easiest way to run inference is to simply use these commands.

For each of the desired configurations, run:

nnUNet_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -t TASK_NAME_OR_ID -m CONFIGURATION --save_npz

Only specify — save_npz if you intend to use ensembling. — save_npz will make the command save the softmax probabilities alongside of the predicted segmentation masks requiring a lot of disk space.

Note: Please select a separate OUTPUT_FOLDER for each configuration. If you interrupted the training, then then rename model_best.model.pkl to model_final_checkpoint.model.pkl and model_best.model to model_final_checkpoint.model for the given fold.

If you wish to run ensembling, you can ensemble the predictions from several configurations with the following command:

nnUNet_ensemble -f FOLDER1 FOLDER2 ... -o OUTPUT_FOLDER -pp POSTPROCESSING_FILE

While trianing, we trained the dataset to learn, White Matter as well as Grey Matter. But for the challenge we only need to predict Gray matter labelled as 1, and everything else is 0. So we convert the labels.

Test images naming convention: [Original test filename]-[team name].nii.gz

Test Results: Download the test folder, and upload with at the following link.

Note: These results are achieved with only 10 epochs on 3D full resolution on a single fold without post-processing, and it is better than results, mentioned by SCGM Challenge website here. Dice Score is better by 0.04 from the best results.