[论文解读] Anatomically-Controllable Medical Image Generation with Segmentation-Guided Diffusion Models

Nicholas Konz, Yuwen Chen+2 more|arXiv (Cornell University)|Feb 7, 2024

Medical Image Segmentation Techniques被引用 4

一句话总结

本文提出一种基于分割引导的扩散模型，通过在采样过程中以部分或完整的解剖分割掩码作为条件，实现解剖学上逼真的医学图像生成。通过一种新颖的掩码擦除训练策略，该模型即使在某些结构被遮蔽的情况下，也能生成高保真、解剖学上合理的图像，在乳腺MRI和CT数据集上，其解剖学真实感和掩码忠实度均优于SOTA模型如ControlNet。

ABSTRACT

Diffusion models have enabled remarkably high-quality medical image generation, yet it is challenging to enforce anatomical constraints in generated images. To this end, we propose a diffusion model-based method that supports anatomically-controllable medical image generation, by following a multi-class anatomical segmentation mask at each sampling step. We additionally introduce a random mask ablation training algorithm to enable conditioning on a selected combination of anatomical constraints while allowing flexibility in other anatomical areas. We compare our method ("SegGuidedDiff") to existing methods on breast MRI and abdominal/neck-to-pelvis CT datasets with a wide range of anatomical objects. Results show that our method reaches a new state-of-the-art in the faithfulness of generated images to input anatomical masks on both datasets, and is on par for general anatomical realism. Finally, our model also enjoys the extra benefit of being able to adjust the anatomical similarity of generated images to real images of choice through interpolation in its latent space. SegGuidedDiff has many applications, including cross-modality translation, and the generation of paired or counterfactual data. Our code is available at https://github.com/mazurowski-lab/segmentation-guided-diffusion.

研究动机与目标

• 解决基于扩散模型的医学图像生成中全局解剖学真实感不足的挑战。
• 实现在图像生成过程中对部分或不完整解剖分割掩码的灵活条件控制。
• 通过训练过程中引入掩码擦除的自监督学习，提升无条件图像生成的质量。
• 支持预配准图像合成与反事实分析等下游应用。
• 在解剖学保真度和输入掩码遵循度方面超越现有掩码条件与无条件生成模型。

提出的方法

• 模型在每个去噪步骤中，基于分割掩码进行图像空间扩散过程的条件控制。
• 采用一种掩码擦除训练策略，使模型暴露于输入掩码中缺失解剖类别所有可能组合的情形。
• 该训练方法使模型能够在未被遮蔽区域推断并生成逼真的解剖结构，同时忠实遵循提供的约束条件。
• 模型在真实医学图像及其对应分割掩码上进行端到端训练。
• 使用辅助分割网络评估生成图像的掩码忠实度。
• 该框架支持条件生成（使用部分掩码）与无条件生成（使用空掩码），后者因自监督学习而实现更高真实感。

Figure 1: Standard diffusion models (right) can fail to create realistic tissue, even if the overall image appears high-quality. From left to right: (1) randomly sampled real breast MRI images, (2) synthetic images created by our segmentation-guided diffusion model, and (3) by the same model but wit

实验结果

研究问题

• RQ1 通过在每个去噪步骤中引入分割掩码引导，扩散模型能否生成更具解剖学真实感的医学图像？
• RQ2 当提供不完整或被擦除的掩码时，模型能否学会在未被遮蔽区域生成合理的解剖结构？
• RQ3 掩码擦除训练是否能提升模型在无条件（空掩码）设置下生成逼真图像的能力？
• RQ4 与SOTA模型如ControlNet相比，该模型在解剖学真实感和掩码忠实度方面的表现如何？
• RQ5 该模型能否用于为下游应用合成配对的、已预配准的医学图像？

主要发现

• 在乳腺MRI数据集上，该模型的掩码忠实度Dice分数达到0.9027，CT器官数据集上为0.8980，显著优于ControlNet的0.3636和0.1132。
• 在相同数据集上，该模型在生成图像与真实图像预测之间的分割重叠Dice分数分别达到0.8593和0.8797，表明具有高度的解剖学保真度。
• 即使在空掩码（无条件生成）条件下，该模型生成的图像也比标准无条件扩散模型更具解剖学真实感，这是由于训练过程中通过掩码擦除实现的自监督学习。
• 在合成数据上训练的分割网络性能几乎与在真实数据上训练的模型相当，Dice差异小于0.04，证实了生成图像的高真实感。
• 当仅对部分结构（如血管和纤维腺体组织）施加约束时，该模型成功生成了多样化且解剖学上合理的图像，同时允许其他结构（如乳房形状）保持变化。
• 作者公开发布了代码库及一对配对的、已预配准的生成乳腺MRI数据集，以支持社区研究。

Figure 3: Breast MRI (top) and CT organ (bottom) images generated with different segmentation-conditional and unconditional models. Samples from conditional models ( blue ) are below their input mask, and unconditional model ( red ) samples are random. For breast MRI, the breast, blood vessel, and f