保姆级教程：nnUnet在2维图像的训练和测试
保姆级教程：nnUnet在2维图像的训练和测试
保姆级教程：nnUnet在2维图像的训练和测试
一、 nnUnet介绍
nnUnet方法源自论文
《Automated Design of Deep Learning Methods for Biomedical Image Segmentation》，来自德国癌症研究中心。
源码地址：https://github.com/MIC-DKFZ/nnUNet。 nnUNet浓缩了医学图像的语义分割领域的大部分知识，并且具备自动为不同任务设计不同训练方案的框架，不需要人工进行调参。
流程如下： 1. 自动对不同的数据属性（图片质量、图片模态、图片大小、体素大小、类别比率等）进行分析，生成独一无二的数据指纹（data fingerprint），表示数据集的关键属性； 2. 制定独一无二的训练方案管道指纹（pipeline fingerprint），表示‘管道’关键的优化设计； 3. 数据预处理和数据增强； 4. 采用2D、3D和3D_Cascaded三个网络分别训练，得出各自的模型（三个网络结构共享一个“管道指纹”，五折交叉验证）； 5. 选择出最优的模型进行推理。
nnUNet在19个国际竞赛中取得最优结果，而且在49个任务的大多数中，都达到了SOTA级别。这些任务包括器官的分割、器官的子结构分割、肿瘤分割、病灶分割等，而且是在多个模态下，例如MRI、CT、EM等。
基本结构就是Unet
下图左边是原数据的标签，右边是nnUNet模型的推理结果。
二、环境配置
首先创建nnUnet2环境并激活，之后安装必要的包
source activate nnUnet2
在终端中运行导出命令，设置环境变量
export nnUNet_raw_data_base="/data/Project/nnUnet/Data/nnUNet_raw"
export nnUNet_preprocessed="/data/Project/nnUnet/Data/nnUNet_preprocessed"
export RESULTS_FOLDER="/data/Project/nnUnet/Data/nnUNet_trained_models"
二、数据配置
我们的2维原始数据长这样~
将2维数据转换为3维数据，其实就是z轴为1的3维数据，具体代码在关注公众号后回复：2DDataProcessTo3D.py即可下载。
几个注意的点： nnUnet的数据格式是固定的，Task002_Heart由Task+ID+数据名组成，imagesTr是训练数据，imagesTs是测试数据，labelsTr是训练数据的标签，数据样本la_003_0000.nii.gz由case样本名字+模态标志+.nii.gz组成，不同的模态用0000/0001/0002/0003区分，我把新的任务ID设置为100。
示例树结构：
nnUNet_raw_data_base/nnUNet_raw_data/Task002_Heart
├── dataset.json
├── imagesTr
│ ├── la_003_0000.nii.gz
│ ├── la_004_0000.nii.gz
│ ├── ...
├── imagesTs
│ ├── la_001_0000.nii.gz
│ ├── la_002_0000.nii.gz
│ ├── ...
└── labelsTr
├── la_003.nii.gz
├── la_004.nii.gz
├── ...
我们的原始2维数据是RGB三通道的，我们可以把RGB三通道的数据看成3个模态，分别提取不同通道的数据，形状转换成（1，width, height），采用SimpleITK保存为3维数据。 imagesTr：
labelsTr：
制作dataset.json，前面的内容模仿给定的格式，按照自己的数据修改一下内容，
dataset.json里的数据名去掉了模态标志，所以每个数据看起来有三个重复名字，在后面读图的时候会自动添加模态标志。
三、预处理
nnUNet_plan_and_preprocess -t 100 -data23d 2 --verify_dataset_integrity
这里我增加了一个设置参数data23d用来指定要处理的是2维数据还是3维数据，在sanity_checks.py中做如下设置，
if data23d == '2':
expected_train_identifiers = np.unique(expected_train_identifiers)
expected_test_identifiers = np.unique(expected_test_identifiers)
print('train num', len(expected_train_identifiers))
print('test num:', len(expected_test_identifiers))
nnU-Net提取数据集指纹（一组特定于数据集的属性，例如图像大小，体素间距，强度信息等）。nnUNet_plan_and_preprocess使用2D U-Net以及所有适用的3D U-Net的预处理数据创建子文件夹。它还将为2D和3D配置创建“计划”文件（带有ending.pkl）。这些文件包含生成的分割管道配置，并且将被nnUNetTrainer读取。
此信息用于创建三个U-Net配置： 一个2D U-Net 一个以全分辨率图像运行的3D U-Net 一个3D U-Net级联
针对我们的2维数据只创建2D U-Net的计划。
四、训练过程
我是在11G的2080ti GPU卡进行训练，分别做5折交叉验证，如：
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task100_PapSmear 0 --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task100_PapSmear 1 --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task100_PapSmear 2 --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task100_PapSmear 3 --npz
CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 Task100_PapSmear 4 --npz
运行过程如下：
中间的网络结构忽略。。。
这样我们的训练就开始了，每个epoch会打印训练损失、验证损失、几个类的平均dice、当前学习率、运行时间，并保存模型，这里的Dice_nulcear Just For papsmear是我自己单独加的一个dice，大家可以忽略也可以自己插入想监测的参数。 默认训练1000个epoch，之后会生成对应fold的文件：
五、确定最佳的U-Net配置
运行完5折交叉验证后，可以确定最佳的配置，下面的100是Task的ID，我们这个任务是100，
nnUNet_find_best_configuration -m 2d -t 100 –strict
运行后会在路径nnUNet_trained_models/nnUNet/ensembles/Task100_PapSmear下生成
也会生成运行推断的方法
nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr nnUNetTrainerV2 -ctr nnUNetTrainerV2CascadeFullRes -m 2d -p nnUNetPlansv2.1 -t Task100_PapSmear
六、运行推断
根据前面生成的推断方法，修改对应输入-i和输出-o文件夹，
nnUNet_predict -i /data/Project/nnUnet/Data/nnUNet_raw/nnUNet_raw_data/Task100_PapSmear/imagesTs/ -o /data/Project/nnUnet/Data/nnUNet_raw/nnUNet_raw_data/Task100_PapSmear/imagesTsPred/ -tr nnUNetTrainerV2 -ctr nnUNetTrainerV2CascadeFullRes -m 2d -p nnUNetPlansv2.1 -t Task100_PapSmear
运行过程如下：
在前面给定的输出文件夹生成五折交叉验证的集成结果
之后重新读取输出结果后，保存为2维图像，即为分割结果
其中一张图像的分割结果如下：
总体而言，nnUnet的性能还是很强劲的，在之前的研究中，我们分别采用过三种方法对该数据集进行分割：
1：Automated Segmentation of Cervical Nuclei in Pap Smear Images using Deformable Multi-path Ensemble Model，链接https://arxiv.org/abs/1812.00527
2：PGU-net+: Progressive Growing of U-net+ for Automated Cervical Nuclei Segmentation，链接https://arxiv.org/abs/1911.01062
3：ASCNet: Adaptive-Scale Convolutional Neural Networks for Multi-Scale Feature Learning，链接https://arxiv.org/abs/1907.03241
分别获得的分割Dice为：0.933、0.926、0.915，nnUnet的分割Dice为0.9624040666779203！！！
最后祝大家新年快乐，2021年冲鸭！！！