nnunet入门之一 (CT图像分割)

MIC-DKFZ/nnUNet
选择Linux环境运行该项目，Windows环境需要更改较多的参数，暂不支持。

安装环境

1. 安装cuda, cudnn，已安装的检测cuda版本

   检测cuda版本：

   • nvcc -v
   • cd /usr/local
   • nvidia-smi（错误方式）

2. 安装对应cuda版本的pytorch

3. 验证

   python -c 'import torch;print(torch.backends.cudnn.version())'
   # 8005, 显示类似此内容
   

   python -c 'import torch;print(torch.__version__)'   
   # 1.10.1+cu111, 显示类似此内容
   

4. 安装nnUNet

   # 安装nnunet
   pip install nnunet
   git clone https://github.com/MIC-DKFZ/nnUNet.git
   cd nnUNet
   pip install -e .
   

   # 设置环境变量
   vim ~/.bashrc
   export nnUNet_raw_data_base="/root/nnUNet_raw_data_base"
   export nnUNet_preprocessed="/root/nnUNet_preprocessed"
   export RESULTS_FOLDER="/root/nnUNet_trained_models"
   source ~/.bashrc
   

   注意：nnUNet_preprocessed的路径要指向SSD上，否则的话，速度会慢很多。

5. 安装hiddenlayer（可选）

   pip install hiddenlayer
   

数据处理

1. 文件夹目录

   └─Task09_Spleen
       │  dataset.json
       │  
       ├─imagesTr
       │      spleen_2.nii.gz
       │      spleen_3.nii.gz
       │      spleen_6.nii.gz
       │      spleen_8.nii.gz
       │      spleen_9.nii.gz
       │      ...
       ├─imagesTs
       │      spleen_1.nii.gz
       │      spleen_7.nii.gz
       │      ...
       └─labelsTr
              spleen_2.nii.gz
              spleen_3.nii.gz
              spleen_6.nii.gz
              spleen_8.nii.gz
              spleen_9.nii.gz
              ...
   

2. json文件信息

   nnUNet/nnunet/dataset_conversion/utils.py里面的函数generate_dataset_json可以生成相应任务的json文件。

   { 
   "name": "Spleen", 
   "description": "Spleen Segmentation",
   "reference": "Memorial Sloan Kettering Cancer Center",
   "licence":"CC-BY-SA 4.0",
   "release":"1.0 06/08/2018",
   "tensorImageSize": "3D",
   "modality": { 
      "0": "CT"
    }, 
    "labels": { 
      "0": "background", 
      "1": "spleen"
    }, 
    "numTraining": 41, 
    "numTest": 20,
    "training":[{"image":"./imagesTr/spleen_2.nii.gz","label":"./labelsTr/spleen_2.nii.gz"},{"image":"./imagesTr/spleen_3.nii.gz","label":"./labelsTr/spleen_3.nii.gz"},{"image":"./imagesTr/spleen_6.nii.gz","label":"./labelsTr/spleen_6.nii.gz"},{"image":"./imagesTr/spleen_8.nii.gz","label":"./labelsTr/spleen_8.nii.gz"},{"image":"./imagesTr/spleen_9.nii.gz","label":"./labelsTr/spleen_9.nii.gz"},...],
    "test":["./imagesTs/spleen_1.nii.gz","./imagesTs/spleen_7.nii.gz",...]
    }
   
   

   注意："modality"一定要是CT，可以是小写。（因为不同的modality，归一化的方式不同）

3. 转换数据

   nnUNet_convert_decathlon_task -i /xxx/Task09_Spleen
   

   转换的数据存在nnUNet_raw_data_base/nnUNet_raw_data/Task009_Spleen，唯一区别就是imagesTr和imagesTs里的文件名加了后缀"_0000"。

   注意：此处Task09_Spleen变为Task009_Spleen。

预处理

# 只进行3d预处理，不进行2d预处理
nnUNet_plan_and_preprocess -t 09 -pl2d None

主要对数据进行预处理，包括Crop, Resample, Normalization等。

训练

如果训练断了，可以在命令行后面加-c继续训练。

1. Lowres训练

   nnUNet_train 3d_lowres nnUNetTrainerV2 09 0
   

2. Fullres训练

   nnUNet_train 3d_fullres nnUNetTrainerV2 09 0
   

3. Cascade训练

   # 先训练完lowres的所有folds
   nnUNet_train 3d_lowres nnUNetTrainerV2 09 0
   nnUNet_train 3d_lowres nnUNetTrainerV2 09 1
   nnUNet_train 3d_lowres nnUNetTrainerV2 09 2
   nnUNet_train 3d_lowres nnUNetTrainerV2 09 3
   nnUNet_train 3d_lowres nnUNetTrainerV2 09 4
   

   nnUNet_train 3d_cascade_fullres nnUNetTrainerV2CascadeFullRes 09 0
   

测试

1. 预测

   # lowres预测
   nnUNet_predict -i /root/nnUNet_raw_data_base/nnUNet_raw_data/Task009_Spleen/imagesTs -o /root/nnUNet_trained_models/nnUNet/3d_lowres/Task009_Spleen/nnUNetTrainerV2__nnUNetPlansv2.1/fold_0/infer -t 09 -m 3d_lowres
   

   # fullres预测
   nnUNet_predict -i /root/nnUNet_raw_data_base/nnUNet_raw_data/Task009_Spleen/imagesTs -o /root/nnUNet_trained_models/nnUNet/3d_fullres/Task009_Spleen/nnUNetTrainerV2__nnUNetPlansv2.1/fold_0/infer -t 09 -m 3d_fullres
   

   # lowres预测
   nnUNet_predict -i /root/nnUNet_raw_data_base/nnUNet_raw_data/Task009_Spleen/imagesTs -o /root/nnUNet_trained_models/nnUNet/3d_cascade_fullres/Task009_Spleen/nnUNetTrainerV2CascadeFullRes__nnUNetPlansv2.1/fold_0/infer -t 09 -m 3d_cascade_fullres
   

2. 评价

   nnUNet_evaluate_folder -ref labelsTs_dir -pred infer_dir -l 1
   

   labelsTs_dir, 测试集的label；infer_dir, 测试集的预测mask。

   如果是四分类，则-l 1 2 3