MedicalSeg-nnUNet使用教程

在医疗图像分割任务中，常常针对不同的数据集设计不同的模型达到较好的性能，但是这些模型在其他数据集上的表现不一定会好。
针对这个问题，MSD（Medical Segmentation Decathlon Challenge）挑战赛给出10个数据集，要求选手评估算法在10个数据上的性能。
nnUNet获得了MSD挑战赛第一名（下图为MSD第一阶段7个数据集算法性能排名）：

相关链接：
paper：nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation
MedicalSeg-nnUNet：https://github.com/PaddlePaddle/PaddleSeg/tree/develop/contrib/MedicalSeg/configs/nnunet/msd_lung
nnUNet：https://github.com/MIC-DKFZ/nnUNet

一、功能模块

nnUNet各功能模块如下图（本项目不会过多介绍其原理，具体内容可参考论文）：

二、MedicalSeg-nnUNet功能模块

MedicalSeg-nnUNet实现的功能：

三、快速体验

按照以下步骤快速体验吧。
PS:2D-UNet和3D-UNet的数据集处理后会大于100GB，训练一个模型的时候，把另一个的数据集删除。

1、准备工作

# PaddleSeg develop分支支持nnUNet（2022年10月25日）
%cd ~/
!git clone https://github.com/PaddlePaddle/PaddleSeg.git
!git pull
%cd ~/PaddleSeg/
!git checkout develop

# 解压数据集
%cd ~/
!mkdir ~/PaddleSeg/contrib/MedicalSeg/msd_lung
!tar -xf ~/data/data125872/Task06_Lung.tar -C ~/PaddleSeg/contrib/MedicalSeg/msd_lung

# 安装依赖
%cd ~/PaddleSeg/contrib/MedicalSeg
!pip install -r requirements.txt
!pip install medpy

# 解压2D-UNet和Cascade UNet-lowres训练好的权重
%cd ~/
!mkdir ~/PaddleSeg/contrib/MedicalSeg/output
!unzip -oq ~/data/data169512/2d_and_cascade_lowres.zip -d ~/PaddleSeg/contrib/MedicalSeg/output/

2、2D-UNet\3D-UNet训练验证

五折训练，对应五个配置文件，初次训练会进行数据预处理，耗时较久。

训练

为了加速训练，使用混合精度训练。nnUNet使用了五折交叉验证的训练方法提升鲁棒性，配置文件的参数不建议修改（默认参数和nnUNet官方repo一致）。
训练命令和paddleseg使用几乎相同，只需要添加–nnunet参数。
注意：不支持多卡训练。

# 2D-UNet 训练
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python train.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold0.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/2d_unet/fold0 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold1.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/2d_unet/fold1 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold2.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/2d_unet/fold2 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold3.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/2d_unet/fold3 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold4.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/2d_unet/fold4 --save_interval 1000 --use_vdl

# 3D-UNet训练
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python train.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold0.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/3d_unet/fold0 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold1.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/3d_unet/fold1 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold2.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/3d_unet/fold2 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold3.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/3d_unet/fold3 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold4.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/3d_unet/fold4 --save_interval 1000 --use_vdl

验证

因为采用五折的训练方法，在验证时，需要五个模型在各自的验证集上验证性能，然后将五个验证集整合到一起验证。
模型参数的选择：训练时开启评估，验证集是经过采样-数据增强后生成的（因为验证一个完整数据会非常耗时间，所以用了近似数据集替代验证集），建议使用最后一次的权重进行验证（和nnUNet一致，也支持选择其他权重）。

# 2D UNet验证-提供权重
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold0.yml --model_path output/2d_unet/fold0/iter_30000/model.pdparams --val_save_folder output/2d_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold1.yml --model_path output/2d_unet/fold1/iter_30000/model.pdparams --val_save_folder output/2d_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold2.yml --model_path output/2d_unet/fold2/iter_30000/model.pdparams --val_save_folder output/2d_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold3.yml --model_path output/2d_unet/fold3/iter_30000/model.pdparams --val_save_folder output/2d_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold4.yml --model_path output/2d_unet/fold4/iter_30000/model.pdparams --val_save_folder output/2d_val --precision fp16

!python nnunet/all_folds_eval.py --gt_dir output/2d_val/gt_niftis --val_pred_dir output/2d_val

# 3D UNet 验证
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold0.yml --model_path output/3d_unet/fold0/iter_30000/model.pdparams --val_save_folder output/3dfullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold1.yml --model_path output/3d_unet/fold1/iter_30000/model.pdparams --val_save_folder output/3dfullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold2.yml --model_path output/3d_unet/fold2/iter_30000/model.pdparams --val_save_folder output/3dfullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold3.yml --model_path output/3d_unet/fold3/iter_30000/model.pdparams --val_save_folder output/3dfullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3d_fullres_msd_lung_fold4.yml --model_path output/3d_unet/fold4/iter_30000/model.pdparams --val_save_folder output/3dfullres_val --precision fp16

!python nnunet/all_folds_eval.py --gt_dir output/3dfullres_val/gt_niftis --val_pred_dir output/3dfullres_val

3、Cascade UNet训练、验证

Cascade UNet分为lowres和fullres 共2个阶段，第一阶段的预测结果是第二阶段的输入。
lowres：第一阶段和2D-UNet以及3D-UNet使用方法一致。
fullres：第二阶段需要第一阶段的预测结果，所以训练第二阶段之前，必须完成第一阶段的【训练-验证】2个步骤，在验证时，命令中加上–predict_next_stage参数。

# Cascade UNet lowres 
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold0.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold0 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold1.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold1 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold2.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold2 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold3.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold3 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold4.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_lowres/fold4 --save_interval 1000 --use_vdl

# Cascade UNet lowres验证，并进行推理，得到下一阶段的输入-提供权重
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold0.yml --model_path output/cascade_lowres/fold0/iter_30000/model.pdparams --val_save_folder output/cascade_lowres_val --precision fp16 --predict_next_stage
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold1.yml --model_path output/cascade_lowres/fold1/iter_30000/model.pdparams --val_save_folder output/cascade_lowres_val --precision fp16 --predict_next_stage
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold2.yml --model_path output/cascade_lowres/fold2/iter_30000/model.pdparams --val_save_folder output/cascade_lowres_val --precision fp16 --predict_next_stage
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold3.yml --model_path output/cascade_lowres/fold3/iter_30000/model.pdparams --val_save_folder output/cascade_lowres_val --precision fp16 --predict_next_stage
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_lowres_msd_lung_fold4.yml --model_path output/cascade_lowres/fold4/iter_30000/model.pdparams --val_save_folder output/cascade_lowres_val --precision fp16 --predict_next_stage

!python nnunet/all_folds_eval.py --gt_dir output/cascade_lowres_val/gt_niftis --val_pred_dir output/cascade_lowres_val

# Cascade UNet fullres 训练
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold0.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_fullres/fold0 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold1.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_fullres/fold1 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold2.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_fullres/fold2 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold3.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_fullres/fold3 --save_interval 1000 --use_vdl
!python train.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold4.yml --log_iters 20 --precision fp16 --nnunet --save_dir output/cascade_fullres/fold4 --save_interval 1000 --use_vdl

# Cascade UNet fullres 验证
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold0.yml --model_path output/cascade_fullres/fold0/iter_30000/model.pdparams --val_save_folder output/cascade_fullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold1.yml --model_path output/cascade_fullres/fold1/iter_30000/model.pdparams --val_save_folder output/cascade_fullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold2.yml --model_path output/cascade_fullres/fold2/iter_30000/model.pdparams --val_save_folder output/cascade_fullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold3.yml --model_path output/cascade_fullres/fold3/iter_30000/model.pdparams --val_save_folder output/cascade_fullres_val --precision fp16
!python nnunet/single_fold_eval.py --config configs/nnunet/msd_lung/nnunet_3dcascade_fullres_msd_lung_fold4.yml --model_path output/cascade_fullres/fold4/iter_30000/model.pdparams --val_save_folder output/cascade_fullres_val --precision fp16

!python nnunet/all_folds_eval.py --gt_dir output/cascade_fullres_val/gt_niftis --val_pred_dir output/cascade_fullres_val

4、Ensemble

对各个模型（2D-UNet、3D-UNet、Cascade UNet）成对Ensemble
为了提升性能，使用模型集成的方法（下面最下面的3个策略）。nnUNet最后会评7个策略的精度，挑选精度最高的提交，分别是：
2D-UNet
3D-UNet
Cascade-UNet-lowres
Cascade-UNet-fullres
2D-UNet + 3D-UNet
2D-UNet + Cascade-UNet-fullres
3D-UNet + Cascade-UNet-fullres

# 2d and cascade lowres ensemble
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python nnunet/ensemble.py --ensemble_folds output/3dfullres_val output/cascade_lowres_val output/2d_val --gt_dir output/cascade_lowres_val/gt_niftis --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_2D.pkl --output_folder output/ensemble

5、预测

训练好动态图模型，使用多折预测的方法进行预测，命令如下。
预测的结果同样支持模型间Ensemble。如果预测的结果需要Ensemble，加上–save_npz参数。
比如想要：2D-UNet + 3D-UNet 的预测，就分别得到2个模型的预测结果，使用上述4的Ensemble命令集成。

# 2D UNet 预测--已提供权重
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python nnunet/predict.py --image_folder msd_lung/Task006_Lung/imagesTs  --output_folder output/nnunet_predict/2d_unet --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_2D.pkl --model_paths output/2d_unet/fold0/iter_30000/model.pdparams output/2d_unet/fold1/iter_30000/model.pdparams output/2d_unet/fold2/iter_30000/model.pdparams output/2d_unet/fold3/iter_30000/model.pdparams output/2d_unet/fold4/iter_30000/model.pdparams --postprocessing_json_path output/2d_val/postprocessing.json --model_type 2d --save_npz

动态图预测示例（测试集008）：

# 3D UNet 预测
%cd ~/PaddleSeg/contrib/MedicalSeg/
!python nnunet/predict.py --image_folder msd_lung/Task006_Lung/imagesTs  --output_folder output/nnunet_predict/3d_unet --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_3D.pkl --model_paths output/3d_unet/fold0/iter_30000/model.pdparams output/3d_unet/fold1/iter_30000/model.pdparams output/3d_unet/fold2/iter_30000/model.pdparams output/3d_unet/fold3/iter_30000/model.pdparams output/3d_unet/fold4/iter_30000/model.pdparams --postprocessing_json_path output/3d_unet/postprocessing.json --model_type 3d --save_npz

# Cascade UNet 预测，先使用lowres预测，在fullres预测
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python nnunet/predict.py --image_folder msd_lung/Task006_Lung/imagesTs  --output_folder output/nnunet_predict/lowres_pred --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_3D.pkl --model_paths output/cascade_lowres/fold0/iter_30000/model.pdparams output/cascade_lowres/fold1/iter_30000/model.pdparams output/cascade_lowres/fold2/iter_30000/model.pdparams output/cascade_lowres/fold3/iter_30000/model.pdparams output/cascade_lowres/fold4/iter_30000/model.pdparams --postprocessing_json_path output/cascade_lowres_val/postprocessing.json --model_type cascade_lowres
!python nnunet/predict.py --image_folder msd_lung/Task006_Lung/imagesTs  --output_folder output/nnunet_predict/fullres_pred --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_3D.pkl --model_paths output/cascade_fullres/fold0/iter_30000/model.pdparams output/cascade_fullres/fold1/iter_30000/model.pdparams output/cascade_fullres/fold2/iter_30000/model.pdparams output/cascade_fullres/fold3/iter_30000/model.pdparams output/cascade_fullres/fold4/iter_30000/model.pdparams --postprocessing_json_path output/cascade_lowres_val/postprocessing.json --model_type cascade_fullres --save_npz --lowres_segmentations output/nnunet_predict/lowres_pred

6、部署

仅以2D UNet为示例，演示模型的部署过程。
1、导出静态图模型（五折训练对应五个模型）；
2、静态图推理（五折训练就对应五折推理）。
PS：静态图推理的模型也支持Ensemble，需要在推理时加上–save_npz参数，Ensemble命令和4一致。

# 静态图导出
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python nnunet/export.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold0.yml --save_dir output/static/2d_unet/fold0 --model_path output/2d_unet/fold0/iter_30000/model.pdparams
!python nnunet/export.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold1.yml --save_dir output/static/2d_unet/fold1 --model_path output/2d_unet/fold1/iter_30000/model.pdparams
!python nnunet/export.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold2.yml --save_dir output/static/2d_unet/fold2 --model_path output/2d_unet/fold2/iter_30000/model.pdparams
!python nnunet/export.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold3.yml --save_dir output/static/2d_unet/fold3 --model_path output/2d_unet/fold3/iter_30000/model.pdparams
!python nnunet/export.py --config configs/nnunet/msd_lung/nnunet_2d_msd_lung_fold4.yml --save_dir output/static/2d_unet/fold4 --model_path output/2d_unet/fold4/iter_30000/model.pdparams

# 静态图多折推理
%cd ~/PaddleSeg/contrib/MedicalSeg/

!python nnunet/infer.py --image_folder msd_lung/Task006_Lung/imagesTs  --output_folder output/nnunet_static/2d_unet --plan_path msd_lung/preprocessed/Task006_Lung/nnUNetPlansv2.1_plans_2D.pkl --model_paths output/static/2d_unet/fold0/model.pdmodel output/static/2d_unet/fold1/model.pdmodel output/static/2d_unet/fold2/model.pdmodel output/static/2d_unet/fold3/model.pdmodel output/static/2d_unet/fold4/model.pdmodel --param_paths output/static/2d_unet/fold0/model.pdiparams output/static/2d_unet/fold1/model.pdiparams output/static/2d_unet/fold2/model.pdiparams output/static/2d_unet/fold3/model.pdiparams output/static/2d_unet/fold4/model.pdiparams --postprocessing_json_path output/2d_unet_val/postprocessing.json --model_type 2d --disable_postprocessing  --save_npz

测试集008静态图预测结果如下：

四、总结

本教程目的是提供nnUNet的使用案例，对nnUNet原理并无过多涉及，请自行阅读相关论文。

请点击此处查看本环境基本用法.

Please click here for more detailed instructions.
此文章为搬运
原项目链接