MedSegDiff：基于扩散模型的医学图像分割

MedSegDiff 是一个基于扩散概率模型（DPM）的框架，用于从医学图像中对器官/组织进行分割和重建。

该算法在我们的论文中有详细阐述：

MedSegDiff：基于扩散概率模型的医学图像分割

MedSegDiff-V2：基于Transformer的扩散式医学图像分割.

快速概览

MedSegDiff-V1	MedSegDiff-V2

最新消息

• [置顶] 加入我们的 Discord，提问并与大家交流。
• 22-11-30。本项目仍在快速更新中。请查看 TODO 列表，了解接下来将发布的内容。
• 22-12-03。修复了 BraTs2020 的一些错误，并添加了示例案例。
• 22-12-15。修复了多 GPU 分布式训练的问题。
• 22-12-16。DPM-Solver ✖️ MedSegDiff 完成 🥳 现在 DPM-Solver 已集成到 MedsegDiff 中。通过设置 --dpm_solver True，即可享受其闪电般的采样速度（1000 步 ❌ 20 步 ⭕️）。
• 22-12-23。修复了 DPM-Solver 中的一些 bug。
• 23-01-31。MedSegDiff-V2 即将发布 🥳 。请先阅读我们的论文 MedSegDiff-V2：基于Transformer的扩散式医学图像分割。
• 23-02-07。优化了 BRATS 数据采样的工作流程，并新增了用于处理原始 3D BRATS 数据的数据加载器。
• 23-02-11。修复了 3D BRATS 数据训练中的 bug，参见 issue 31。
• 23-03-04。论文 MedSegDiff：基于扩散概率模型的医学图像分割 已被 MIDL 2023 正式接受 🥳
• 23-04-11。基于 V2 框架的新版本已发布 🥳。相比之前版本，它更加准确、稳定且具有更好的领域适应性，同时不会占用过多资源。我们还修复了一些小问题，比如 requirement.txt 和 isic csv 文件。非常感谢所有报告问题的用户，你们的帮助对我们至关重要 🤗。顺便说一下，现在默认运行的是新版本。如果想使用旧版本，请添加 --version 1。
• 23-04-12。为 isic 数据集添加了一个简单的评估文件（script/segmentation_env）。使用方法：python scripts/segmentation_env.py --inp_pth *保存预测图像的文件夹* --out_pth *保存真实标签图像的文件夹*
• 23-12-05。论文 MedSegDiff-V2：基于Transformer的扩散式医学图像分割 已被 AAAI 2024 正式接受 🥳
• 24-10-29。MedSegDiff-V2 被评选为 AAAI-24 最具影响力论文 🥳

系统要求

pip install -r requirement.txt

示例案例

皮肤图像中的黑色素瘤分割

1. 从 https://challenge.isic-archive.com/data/ 下载 ISIC 数据集。您的数据集文件夹应位于 "data" 目录下，结构如下：

data
|   ----ISIC
|       ----Test
|       |   |   ISBI2016_ISIC_Part1_Test_GroundTruth.csv
|       |   |   
|       |   ----ISBI2016_ISIC_Part1_Test_Data
|       |   |       ISIC_0000003.jpg
|       |   |       .....
|       |   |
|       |   ----ISBI2016_ISIC_Part1_Test_GroundTruth
|       |           ISIC_0000003_Segmentation.png
|       |   |       .....
|       |           
|       ----Train
|           |   ISBI2016_ISIC_Part1_Training_GroundTruth.csv
|           |   
|           ----ISBI2016_ISIC_Part1_Training_Data
|           |       ISIC_0000000.jpg
|           |       .....
|           |       
|           ----ISBI2016_ISIC_Part1_Training_GroundTruth
|           |       ISIC_0000000_Segmentation.png
|           |       .....

2. 训练时，运行：python scripts/segmentation_train.py --data_name ISIC --data_dir *输入数据路径* --out_dir *输出数据路径* --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 1e-4 --batch_size 8

3. 采样时，运行：python scripts/segmentation_sample.py --data_name ISIC --data_dir *输入数据路径* --out_dir *输出数据路径* --model_path *保存的模型* --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --num_ensemble 5

4. 评估时，运行：python scripts/segmentation_env.py --inp_pth *保存预测图像的文件夹* --out_pth *保存真实标签图像的文件夹*

默认情况下，样本将保存在 ./results/

MRI 图像中的脑肿瘤分割

1. 从 https://www.med.upenn.edu/cbica/brats2020/data.html 下载 BRATS2020 数据集。您的数据集文件夹应如下所示：

data
└───training
│   └───slice0001
│       │   brats_train_001_t1_123_w.nii.gz
│       │   brats_train_001_t2_123_w.nii.gz
│       │   brats_train_001_flair_123_w.nii.gz
│       │   brats_train_001_t1ce_123_w.nii.gz
│       │   brats_train_001_seg_123_w.nii.gz
│   └───slice0002
│       │  ...
└───testing
│   └───slice1000
│       │  ...
│   └───slice1001
│       │  ...

2. 训练时，运行：python scripts/segmentation_train.py --data_dir (您放置数据文件夹的位置)/data/training --out_dir 输出数据路径 --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 1e-4 --batch_size 8

3. 采样时，运行：python scripts/segmentation_sample.py --data_dir (您放置数据文件夹的位置)/data/testing --out_dir 输出数据路径 --model_path 保存的模型 --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --num_ensemble 5

其他示例

...

在您自己的数据集上运行

在其他数据集上运行 MedSegDiff 非常简单。只需按照 ./guided_diffusion/isicloader.py 或 ./guided_diffusion/bratsloader.py 的格式编写另一个数据加载器文件即可。如果您遇到任何问题，欢迎提交 Issue。如果您能贡献自己的数据集扩展，我们将不胜感激。与自然图像不同，医学图像因任务不同而差异很大。要扩大方法的泛化能力，需要大家共同努力。

超参数与训练建议

要训练出一个优秀的模型，即论文中的 MedSegDiff-B，可将模型超参数设置为：

--image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 

扩散超参数设置为：

--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False

为了加快采样速度：

--diffusion_steps 50 --dpm_solver True 

多 GPU 运行时：

--multi-gpu 0,1,2（例如）

训练超参数设置为：

--lr 5e-5 --batch_size 8

并在采样时设置 --num_ensemble 5。

在大多数数据集上，训练约 10 万步即可收敛。需要注意的是，尽管在后期的许多步骤中损失不再下降，但结果的质量仍在不断提升。这种现象在其他 DPM 应用中也有观察到，比如图像生成。希望有聪明人能告诉我这是为什么🥲。

我很快会发布其在较小批量下的性能表现（适合在 24GB 显存的 GPU 上运行），以便进行对比🤗。

要充分发挥其潜力的配置（MedSegDiff++）是：

--image_size 256 --num_channels 512 --class_cond False --num_res_blocks 12 --num_heads 8 --learn_sigma True --use_scale_shift_norm True --attention_resolutions 24 

然后使用批量 --batch_size 64 训练，并以集合数 --num_ensemble 25 进行采样。

待办事项列表

• 修复 BRATS 中的 bug。添加 BRATS 示例。
• 发布 REFUGE 和 DDIT 数据加载器及示例。
• 使用 DPM 求解器加速采样。
• 深度推理。
• 修复多 GPU 并行中的 bug。
• 在训练过程中进行采样并可视化。
• 发布预处理和后处理流程。
• 发布评估代码。
• 部署至 HuggingFace。
• 完成配置。

感谢

本代码大量借鉴了以下项目：openai/improved-diffusion、WuJunde/MrPrism、WuJunde/DiagnosisFirst、LuChengTHU/dpm-solver、JuliaWolleb/Diffusion-based-Segmentation、hojonathanho/diffusion、guided-diffusion、bigmb/Unet-Segmentation-Pytorch-Nest-of-Unets、nnUnet、lucidrains/vit-pytorch。

引用

请引用以下文献：

@inproceedings{wu2023medsegdiff,
  title={MedSegDiff: 医学图像分割的扩散概率模型},
  author={吴俊德、RAO、方慧慧、张宇、杨叶辉、熊浩义、刘慧英、徐延武},
  booktitle={深度学习在医学影像中的应用},
  year={2023}
}

@article{wu2023medsegdiff,
  title={MedSegDiff-V2：基于 Transformer 的扩散式医学图像分割},
  author={吴俊德、季伟、傅华珠、许敏、金月明、徐延武},
  journal={arXiv 预印本 arXiv:2301.11798},
  year={2023}
}

MedSegDiff 快速上手指南

MedSegDiff 是一个基于扩散概率模型（DPM）的框架，专为医学图像中的器官/组织分割与重建而设计。支持 V1 基础版及最新的 V2（引入 Transformer 架构，精度更高、更稳定）。

环境准备

• 系统要求：Linux 或 macOS（推荐 Linux），需配备 NVIDIA GPU 以进行训练和加速采样。
• 前置依赖：
  • Python 3.8+
  • PyTorch (兼容版本请参考 requirement.txt)
  • CUDA 工具包
• 数据集准备：
  • 皮肤癌分割 (ISIC)：从 ISIC Archive 下载，目录结构需包含 Train 和 Test 文件夹及对应的 CSV 标签文件。
  • 脑肿瘤分割 (BraTS)：从 BraTS2020 下载，需整理为 .nii.gz 格式的多模态 MRI 数据。

安装步骤

1. 克隆项目代码：

   git clone https://github.com/WuJunde/MedSegDiff.git
   cd MedSegDiff
   

2. 安装依赖包（建议配置国内镜像源加速）：

   pip install -r requirement.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下以 ISIC 皮肤黑色素瘤分割 为例，展示最核心的训练、采样与评估流程。请确保已按 README 要求整理好 data/ISIC 目录结构。

1. 模型训练

运行以下命令开始训练。默认使用最新发布的 V2 框架（更精准、资源占用合理）。若需使用旧版 V1，请添加 --version 1 参数。

python scripts/segmentation_train.py \
  --data_name ISIC \
  --data_dir ./data/ISIC \
  --out_dir ./outputs/isic_model \
  --image_size 256 \
  --num_channels 128 \
  --class_cond False \
  --num_res_blocks 2 \
  --num_heads 1 \
  --learn_sigma True \
  --use_scale_shift_norm False \
  --attention_resolutions 16 \
  --diffusion_steps 1000 \
  --noise_schedule linear \
  --rescale_learned_sigmas False \
  --rescale_timesteps False \
  --lr 1e-4 \
  --batch_size 8

提示：通常训练约 100,000 步即可收敛。如需多卡训练，添加 --multi-gpu 0,1,2,3。

2. 图像采样（推理）

使用训练好的模型进行分割预测。推荐使用 DPM-Solver 将采样步数从 1000 步大幅减少至 20-50 步，显著提升速度。

python scripts/segmentation_sample.py \
  --data_name ISIC \
  --data_dir ./data/ISIC/Test \
  --out_dir ./results/isic_pred \
  --model_path ./outputs/isic_model/model_latest.pt \
  --image_size 256 \
  --num_channels 128 \
  --class_cond False \
  --num_res_blocks 2 \
  --num_heads 1 \
  --learn_sigma True \
  --use_scale_shift_norm False \
  --attention_resolutions 16 \
  --diffusion_steps 50 \
  --noise_schedule linear \
  --rescale_learned_sigmas False \
  --rescale_timesteps False \
  --dpm_solver True \
  --num_ensemble 5

参数说明：

• --dpm_solver True：启用快速采样算法。
• --num_ensemble 5：集成采样次数，增加此值可提升结果稳定性。

3. 结果评估

计算预测结果与真实标签（Ground Truth）之间的指标（如 Dice Score 等）。

python scripts/segmentation_env.py \
  --inp_pth ./results/isic_pred \
  --out_pth ./data/ISIC/Test/ISBI2016_ISIC_Part1_Test_GroundTruth

---

其他数据集：若要使用 BraTS 或其他自定义数据集，只需参照 guided_diffusion/isicloader.py 编写新的数据加载器，并调整上述命令中的 --data_name 和数据路径即可。