[论文解读] DiffuseExpand: Expanding dataset for 2D medical image segmentation using diffusion models

Shitong Shao, Xiaohan Yuan+4 more|arXiv (Cornell University)|Apr 26, 2023

Generative Adversarial Networks and Image Synthesis被引用 8

一句话总结

DiffuseExpand 利用扩散概率模型生成高质量、成对的图像-掩模样本用于2D医学图像分割，可实现高效的数据扩增并提升分割性能。它引入一个四阶段管线：掩模优先的合成、分类器引导、加速采样和最终质量筛选。

ABSTRACT

Dataset expansion can effectively alleviate the problem of data scarcity for medical image segmentation, due to privacy concerns and labeling difficulties. However, existing expansion algorithms still face great challenges due to their inability of guaranteeing the diversity of synthesized images with paired segmentation masks. In recent years, Diffusion Probabilistic Models (DPMs) have shown powerful image synthesis performance, even better than Generative Adversarial Networks. Based on this insight, we propose an approach called DiffuseExpand for expanding datasets for 2D medical image segmentation using DPM, which first samples a variety of masks from Gaussian noise to ensure the diversity, and then synthesizes images to ensure the alignment of images and masks. After that, DiffuseExpand chooses high-quality samples to further enhance the effectiveness of data expansion. Our comparison and ablation experiments on COVID-19 and CGMH Pelvis datasets demonstrate the effectiveness of DiffuseExpand. Our code is released at https://github.com/shaoshitong/DiffuseExpand.

研究动机与目标

• 解决由于隐私和标注限制导致的医学图像分割数据稀缺问题。
• 开发基于扩散模型的管线以生成高多样性和对齐的成对图像与掩模样本。
• 通过阶段性条件控制、分类器引导和样本筛选提升数据扩增质量。
• 在COVID-19胸片和CGMH骨盆数据集上证明有效性。
• 提供开源代码，并通过消融研究和对比验证改进。

提出的方法

• 阶段I：在不使用分类器引导的情况下对预训练的扩散概率模型（DPM）进行微调，以将分布映射到医学图像。
• 阶段II：训练分割器以实现分类器引导，并通过 Dice 和 BCE 损失评估图像-掩模对齐。
• 阶段III：使用 DPM Solver++ 在多种条件下合成图像-掩模对，基于分类器引导实现条件采样。
• 阶段IV：使用训练好的网络和 Dice 阈值筛选，丢弃质量较差的样本，确保扩增数据的高质量。
• 引入针对条件采样的连续时间概率流 ODE 形式，并利用带有温度参数 tau 的分类器引导来控制梯度幅度。
• 提供消融研究以比较无分类器引导与分类器引导采样，并评估 tau/s 对多样性和 Dice 性能的影响。

实验结果

研究问题

• RQ1 扩散模型是否能够为2D医学图像分割生成多样且对齐良好的图像-掩模对？
• RQ2 分阶段条件控制和分类器引导是否提高扩增分割数据集的质量和实用性？
• RQ3 样本筛选（阶段IV）对 Dice 分数和样本多样性在数据扩增中的影响？
• RQ4 在医学成像领域，基于扩散的扩增方法与基于 GAN 的或基于增强的做法相比表现如何？

主要发现

• DiffuseExpand 能产生高保真且多样的图像-掩模对，提升在测试数据集上的 Dice 得分。
• 在训练好的分割器的分类器引导配合阶段IV筛选下，图像与掩模的对齐性得到增强。
• 在类似场景下，使用 tau 为条件引导的方式比以梯度缩放值 s 的方式得到更高的 Dice 分数。
• 在测试设置下，基于扩散的扩增方法在 Dice 基于的分割性能上优于 Synthmed 和 XLsor。
• 阶段IV 筛选降低样本多样性但提升整体质量，从而有助于下游分割性能提升。