QuickStart of nnUNet

Posted on In

nnUNet

Environment

  1. Use Linux/MacOS. WSL2 is recommended for Windows.

  2. create virtual environment, anaconda is recommended.

  3. install jupyter notebook, check here

  4. install pytorch>=2.0, check pytorch

Prepare nnUNet

  1. clone and install nnunet

    1
    2
    3
    git clone git@github.com:MIC-DKFZ/nnUNet.git --depth 1
    cd nnUNet/
    pip install -e . # install as develop mode

  2. clone and install hiddenlayer (visualize model topology graph)

    1
    pip install --upgrade git+https://github.com/FabianIsensee/hiddenlayer.git

  3. prepare dataset

    1. create data directory

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      ./nnUNet_raw_data_base/
      ├── nnUNet_preprocessed
      ├── nnUNet_raw
      │   ├── nnUNet_cropped_data
      │   └── nnUNet_raw_data
      │       └── Dataset001_teeth
      │           ├── dataset.json
      │           ├── imagesTr
      │           ├── imagesTs
      |           ├── inferTs
      │           └── labelsTr
      └── nnUNet_results

      dataset.json

      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      {
          "channel_names":
          {
              "0": "CBCT"
          },
          "labels": 
          {
              "background": 0,
              "Teeth": 1
          },
          "numTraining": 12,
          "file_ending": ".nii.gz"
      }

    2. put your training data to imagesTr/labelsTr; imagesTr for testing

    3. remove postfix of labels

      1
      2
      3
      import os
      for file in os.listdir('.'):
          os.renames(file, file.replace('_0000.nii', '.nii'))

    4. modify the nnUNet_raw, nnUNet_preprocessed, nnUNet_results to your own path in paths.py

  4. dataset preprocessing: nnUNetv2_plan_and_preprocess -d <dataset_id> --verify_dataset_integrity

Train

  1. training: nnUNetv2_train <dataset_id> <UNET_CONFIGURATION> <FOLD>
    1. if you do not know how to set <UNET_CONFIGURATION>, just run all of them 2d, 3d_fullres, 3d_lowres, 3d_cascade_lowres and let nnUnet to decide.
    2. FOLD specifies which fold of the 5-fold-cross-validation is trained. use all if you do not want to use FOLD.
    3. add --npz flag to enable nnUNetv2_find_best_configuration later.

Inference

1
2
3
4
5
export INPUT=/lustre/home/acct-laurence/laurence-user1/wuchen/nnUNet/nnUNet_raw_data_base/nnUNet_raw_data/Task001_teeth/imagesTs

export OUTPUT=/lustre/home/acct-laurence/laurence-user1/wuchen/nnUNet/nnUNet_raw_data_base/nnUNet_raw_data/Task001_teeth/inferTs

nnUNetv2_predict -i $INPUT -o $OUTPUT -c 3d_fullres -d 001 -f all