MedSegDiff: Medical Image Segmentation withDiffusion Probabilistic Model 论文阅读

论文地址:[2211.00611] MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model (arxiv.org)

代码:GitHub - KidsWithTokens/MedSegDiff: Medical Image Segmentation with Diffusion Model

摘要。扩散概率模型(Diffusion probabilistic model, DPM)是近年来计算机视觉研究的热点之一。它在Imagen、Latent Diffusion Models和Stable Diffusion等图像生成应用中表现出了令人印象深刻的生成能力，引起了社区的广泛讨论。最近的许多研究还发现，它在许多其他视觉任务中也很有用，比如图像去模糊、超分辨率和异常检测。受DPM成功的启发，我们提出了第一个基于DPM的一般医学图像分割模型，我们将其命名为MedSegDiff。为了增强DPM在医学图像分割中的分步区域注意力，我们提出了动态条件编码，该编码为每个采样步建立状态自适应条件。我们进一步提出Feature Frequency Parser (FF-Parser)来消除高频噪声分量在此过程中的负面影响。我们在三种不同图像模式的医学分割任务上验证了MedSegDiff，即眼底图像的视杯分割、MRI图像的脑肿瘤分割和超声图像的甲状腺结节分割。实验结果表明，MedSegDiff算法在性能上明显优于SOTA算法，表明了该模型的泛化性和有效性。我们的代码发布在https://github.com/WuJunde/MedSegDiff。

1 简介

医学图像分割是将医学图像划分为有意义的区域的过程。分割是许多医学图像分析应用（例如诊断、手术规划和图像引导手术）中的基本步骤。这很重要，因为它可以让医生和其他医疗专业人员更好地了解他们正在查看的内容。它还可以更轻松地比较图像并跟踪随时间的变化。近年来，人们对自动医学图像分割方法越来越感兴趣。这些方法有可能减少手动分割所需的时间和精力，并提高结果的一致性和准确性。随着深度学习技术的发展，越来越多的研究成功地将基于神经网络（NN）的模型应用于医学图像分割任务，从流行的卷积神经网络（CNN）[11]到最近的视觉变换器（ViT） [3,22,12,28]。

最近，扩散概率模型（DPM）[9]作为一类强大的生成模型[27]而受到欢迎，它能够生成具有高多样性和合成质量的图像。最近的大型扩散模型，例如 DALLE2[17]、Imagen[19] 和 Stable Diffusion[18] 已显示出令人难以置信的生成能力。扩散模型最初应用于没有绝对真相的领域。然而，最近的研究表明，它对于地面实况唯一的问题也有效，例如超分辨率[20]和去模糊[24]。

受到 DPM 最近成功的启发，我们为医学图像分割任务设计了一种独特的基于 DPM 的分割模型。据我们所知，我们是第一个在不同图像模态的一般医学图像分割背景下提出基于DPM的模型。我们注意到，在医学图像分割任务中，病变/器官通常是不明确的并且难以与背景区分开。在这种情况下，自适应校准过程是获得精细结果的关键。遵循这种思维方式，我们提出基于普通 DPM 的动态条件编码来设计所提出的模型，名为 MedSegDiff。请注意，在迭代采样过程中，MedSegDiff 以图像先验条件每个步骤，以便从中学习分割图。为了实现自适应区域注意，我们将当前步骤的分割图集成到每一步的图像先验编码中。具体实现是在特征层面上以多尺度的方式将当前步骤的分割掩模与先验图像进行融合。这样，损坏的当前步掩模有助于动态增强条件特征，从而提高重建精度。为了消除此过程中损坏的给定掩模中的高频噪声，我们进一步提出特征频率解析器（FF-Parser）来过滤傅立叶空间中的特征。每个跳跃连接路径上都采用 FF-Parser 以进行多尺度集成。我们在三种不同的医学分割任务（视杯分割、脑肿瘤分割和甲状腺结节分割）上验证了 MedSegDiff。这些任务的图像有不同的模态，分别是眼底图像、脑部CT图像、超声图像。 MedSegDiff 在不同模式的所有三个任务上都优于之前的 SOTA，这表明了所提出方法的泛化性和有效性。简而言之，本文的贡献是： – 首次提出基于 DPM 的通用医学图像分割模型。

– 提出了动态条件编码策略以进行逐步关注。

– 提出FF-Parser来消除高频成分的负面影响。

– 在具有不同图像模式的三种不同医学分割任务上的 SOTA 性能。

2 方法

我们根据[9]中提到的扩散模型设计模型。扩散模型是由两个阶段组成的生成模型，即前向扩散阶段和反向扩散阶段。在前向过程中，分割标签x0通过一系列步骤T逐渐添加高斯噪声。在反向过程中，训练神经网络通过反转噪声过程来恢复原始数据，可以表示为：

其中 θ 是逆向过程参数。从高斯噪声开始，pθ(xT ) = N (xT ; 0; In×n)，其中 I 是原始图像，逆过程将潜变量分布 pθ(xT ) 转换为数据分布 pθ(x0)。与正向过程对称，逆向过程逐步恢复噪声图像以获得最终清晰的分割。

遵循DPM的标准实现，我们采用UNet作为学习网络。如图1所示。为了实现分割，我们以原始图像先验条件为步长估计函数，可以表示为：

其中 EI t 是条件特征嵌入，在我们的例子中是原始图像嵌入，Ex t 是当前步骤的分割图特征嵌入。将这两个分量相加并发送到 UNet 解码器 D 进行重建。

步骤索引 t 与添加的嵌入和解码器功能集成。

在每一个中，它都使用共享的学习查找表进行嵌入，如下[9]。

2.1 动态条件编码

在大多数条件 DPM 中，条件先验将是唯一的给定信息。然而，医学图像分割因其模糊目标而臭名昭著。病变或组织通常很难与其背景区分开来。

低对比度图像模式，例如 MRI 或超声图像，会使情况变得更糟。仅给出静态图像 I 作为每一步的条件将很难学习。为了解决这个问题，我们为每个步骤提出了动态条件编码。我们注意到，一方面，原始图像包含准确的分割目标信息，但很难与背景区分开来；另一方面，当前步骤的分割图包含增强的目标区域，但不准确。这促使我们整合当前步的分割信息xt转化为有条件的原始图像编码以进行互补。

具体来说，我们是在功能层面实现集成。在原始图像编码器中，我们使用当前步骤的编码特征增强其中间特征。条件特征图mk I的每个尺度与具有相同形状的xt编码特征mk x 融合，k是层的索引。融合是通过类似注意力的机制A来实现的。具体来说，两个特征图首先应用层归一化并相乘以获得亲和力图。然后我们将亲和图与条件编码特征相乘来增强注意力区域，即：

其中 ⊗ 表示逐元素乘法，LN 表示层归一化。

该操作应用于中间两个阶段，其中每个阶段都是按照 ResNet34 实现的卷积阶段。这样的策略有助于 MedSegDiff 动态定位和校准分割。尽管该策略很有效，但另一个具体问题是集成 xt 嵌入会引起额外的高频噪声。为了解决这个问题，我们提出 FF-Parser 来约束特征中的高频成分。

2.2 FF解析器

我们以特征集成的路径方式连接FF解析器。它的作用是约束xt特征中与噪声相关的成分。我们的主要想法是学习应用于傅立叶空间特征的参数化注意力（权重）图。给定解码器特征图 m 2 R H×W×C ，我们首先沿空间维度执行 2D FFT（快速傅里叶变换），我们可以将其表示为：

FF-Parser 可以被视为频率滤波器的可学习版本，广泛应用于数字图像处理[16]。与空间注意力不同，它全局调整特定频率的分量。因此可以学习约束自适应积分的高频分量。

2.3 训练和架构

MedSegDiff 按照 DPM [9] 的标准流程进行训练。具体来说，损失可以表示为：

在每次迭代中，将对原始图像 Ii 和分割标签 Si 进行随机采样以进行训练。迭代次数是从均匀分布和高斯分布中采样的。

MedSegDiff 的主要架构是改进的 ResUNet[26]，我们使用 UNet 解码器后面的 ResNet 编码器来实现它。详细的网络设置如下[14]。 I 和 xt 使用两个单独的编码器进行编码。

编码器由三个卷积级组成。每个阶段包含几个残差块。每个阶段的残差块数量遵循ResNet34。每个残差块由两个卷积块组成，每个卷积块由群范数和SiLU[5]活动层和一个卷积层组成。

残差块通过线性层、SiLU 激活和另一个线性层接收时间嵌入。然后将结果添加到第一个卷积块的输出中。将获得的EI和Ext相加并发送到最后的编码阶段。连接标准卷积解码器来预测最终结果。

3 实验

3.1 数据集

我们使用不同的图像模态对三种不同的医疗任务进行了实验，即眼底图像的视杯分割、MRI 图像的脑肿瘤分割和超声图像的甲状腺结节分割。青光眼、甲状腺癌和黑色素瘤诊断的实验在REFUGE-2数据集[6]、BraTs-2021数据集[2]和DDTI数据集[15]上进行，分别包含1200、2000、8046个样本。这些数据集是公开的，带有分割和诊断标签。

训练/验证/测试集按照数据集的默认设置进行分割。

3.2 实施细节

我们分别对模型的巨型、大型、基本和小型变体 MedSegDiff++、MedSegDiff-L、MedSegDiff-B 和 MedSegDiff-S 进行实验。

在 MedSegDiff-S、MedSegDiff-B、MedSegDiff-L、MedSegDiff++ 中，我们分别使用 UNet 进行 4x、5x、6x、6x 下采样。在实验中，我们采用 100 个扩散步骤进行推理，这比之前的大多数研究要小得多[9,14]。所有实验均使用 PyTorch 平台实现，并在 4 个具有 24GB 内存的 Tesla P40 GPU 上进行训练/测试（MedSegDiff++ 和 MedSegDiff-L 除外）。所有图像统一调整为256×256像素的尺寸。使用 AdamW[13] 优化器以端到端的方式训练网络。 MedSegDiff-B 和 MedSegDiff-S 使用 32 个批量大小进行训练，MedSegDiff-L 和 MedSegDiff++ 使用 64 个批量大小进行训练。学习率最初设置为 1 × 10−4 。所有模型在推理中都设置了 25 次集成。我们使用STAPLE[23]算法来融合不同的样本。基于扩散的竞争对手 EnsemDiff[25] 使用相同的设置进行复制，以进行公平比较。

3.3 主要结果

我们将针对三个特定任务提出的 SOTA 分割方法和通用医学图像分割方法进行比较。主要结果如表1所示。表中，ResUnet[26]和BEAL[21]提出用于视盘/视杯分割，TransBTS[22]和EnsemDiff[25]提出用于脑肿瘤分割，MTSeg[7]和UltraUNet [4]提出用于甲状腺结节分割，CENet[8]、MRNet[11]、SegNet[1]、nnUNet[10]和TransUNet[3]提出用于一般医学图像分割。我们通过 Dice 分数和 IoU 来评估分割性能。

在表1中，我们比较了使用各种网络架构实现的方法，包括CNN（ResUNet、BEAL、nnUNet、SegNet）、视觉变换器（TransBTS、TransUNet）和DPM（EnsemDiff）。我们可以看到先进的网络架构通常会获得更好的结果。例如，在视杯分割中，基于 ViT 的通用分割方法：TransUNet 甚至比基于 CNN 的任务方法：BEAL 更好。在脑肿瘤分割方面，最近提出的基于 DPM 的分割方法 EnsemDiff 优于之前所有基于 ViT 的竞争对手，即 TransBTS 和 TransUNet。 MedSegDiff不仅采用了最近成功的DPM，而且还专门针对一般医学图像分割任务设计了适当的策略。我们可以看到 MedSegDiff 在三个不同的任务上优于所有其他方法，这显示了对不同医学分割任务和不同图像模态的泛化。与专门针对脑肿瘤分割提出的基于 DPM 的模型（即 EnsemDiff）相比，它在 Dice 上提高了 2.3％，在 IoU 上提高了 2.4％，这表明了我们独特的技术（即动态调节和 FF-Parser）的有效性。

图 3 显示了我们的 MedSegDiff 和其他 SOTA 方法生成的几个典型示例。可以看到目标病灶/组织都是不明确的使图像难以被人眼识别。与这些计算机辅助方法相比，很明显，该方法生成的分割图比其他方法更准确，特别是对于模糊区域。受益于 DPM 以及所提出的动态调节和 FF-Parser，它可以更好地定位和校准低对比度或模糊图像的分割。

3.4 消融研究

我们进行了全面的消融研究，以验证所提出的动态调节和 FF-Parser 的有效性。结果如表2所示，其中Dy-Cond表示动态调节。我们通过 Dice 分数 (%) 来评估所有三项任务的表现。从表中我们可以看到 Dy-Cond 比普通 DPM 有了相当大的改进。在区域定位很重要的情况下，即视杯分割，它提高了 2.1%。在图像对比度较低的情况下，如脑肿瘤和甲状腺结节分割，分别提高了1.6%和1.8%。它表明 Dy-Cond 对于这两种情况都是 DPM 上普遍有效的策略。在 Dy-Cond 上建立的 FF-Parser 减轻了高频噪声，从而进一步优化了分割结果。它帮助 MedSegDiff 进一步提高近 1% 的性能，并在所有三项任务上取得最佳成绩。

4。结论

在本文中，我们提供了一种基于 DPM 的通用医学图像分割方案，名为 MedSegDiff。我们提出了两种新颖的技术来保证其性能，即动态条件编码和 FF-Parser。

在具有不同图像模态的三个医学图像分割任务上进行了比较实验，这表明我们的模型优于之前的 SOTA。作为一般医学图像分割中的第一个 DPM 应用，我们相信 MedSegDiff 将成为未来研究的重要基准。

References
1. Badrinarayanan, V., Kendall, A., Cipolla, R.: Segnet: A deep convolutional encoder-decoder architecture for image segmentation. IEEE transactions on pattern analysis and machine intelligence 39(12), 2481–2495 (2017)
2. Baid, U., Ghodasara, S., Mohan, S., Bilello, M., Calabrese, E., Colak, E., Farahani, K., Kalpathy-Cramer, J., Kitamura, F.C., Pati, S., et al.: The rsna-asnr-miccai brats 2021 benchmark on brain tumor segmentation and radiogenomic classification. arXiv preprint arXiv:2107.02314 (2021)
3. Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., Zhou, Y.: Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04306 (2021)
4. Chu, C., Zheng, J., Zhou, Y.: Ultrasonic thyroid nodule detection method based on u-net network. Computer Methods and Programs in Biomedicine 199, 105906 (2021)
5. Elfwing, S., Uchibe, E., Doya, K.: Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks 107, 3–11 (2018)
6. Fang, H., Li, F., Fu, H., Sun, X., Cao, X., Son, J., Yu, S., Zhang, M., Yuan, C.,Bian, C., et al.: Refuge2 challenge: Treasure for multi-domain learning in glaucoma assessment. arXiv preprint arXiv:2202.08994 (2022)
7. Gong, H., Chen, G., Wang, R., Xie, X., Mao, M., Yu, Y., Chen, F., Li, G.: Multitask learning for thyroid nodule segmentation with thyroid region prior. In: 2021 IEEE 18th InternationalSymposium on Biomedical Imaging (ISBI). pp. 257–261. IEEE (2021)
8. Gu, Z., Cheng, J., Fu, H., Zhou, K., Hao, H., Zhao, Y., Zhang, T., Gao, S., Liu J.: Ce-net: Context encoder network for 2d medical image segmentation. IEEE transactions on medical imaging 38(10), 2281–2292 (2019)
9. Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33, 6840–6851 (2020)
10. Isensee, F., Jaeger, P.F., Kohl, S.A., Petersen, J., Maier-Hein, K.H.: nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods 18(2), 203–211 (2021)
11. Ji, W., Yu, S., Wu, J., Ma, K., Bian, C., Bi, Q., Li, J., Liu, H., Cheng, L., Zheng, Y.: Learning calibrated medical image segmentation via multi-rater agreement modeling. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 12341–12351 (2021)
12. Liu, C., Zhao, R., Shi, Z.: Remote-sensing image captioning based on multilayer aggregated transformer. IEEE Geoscience and Remote Sensing Letters 19, 1–5 (2022)
13. Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 (2017)
14. Nichol, A.Q., Dhariwal, P.: Improved denoising diffusion probabilistic models. In: International Conference on Machine Learning. pp. 8162–8171. PMLR (2021)
15. Pedraza, L., Vargas, C., Narváez, F., Durán, O., Muñoz, E., Romero, E.: An open access thyroid ultrasound image database. In: 10th International symposium on medical information processing and analysis. vol. 9287, pp. 188–193. SPIE (2015)
16. Pitas, I.: Digital image processing algorithms and applications. John Wiley & Sons (2000)

17. Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., Chen, M.: Hierarchical textconditional image generation with clip latents. arXiv preprint arXiv:2204.06125 (2022)
18. Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 10684–10695 (2022)
19. Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E., Ghasemipour, S.K.S., Ayan, B.K., Mahdavi, S.S., Lopes, R.G., et al.: Photorealistic textto-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 (2022)
20. Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D.J., Norouzi, M.: Image superresolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence (2022)
21. Wang, S., Yu, L., Li, K., Yang, X., Fu, C.W., Heng, P.A.: Boundary and entropydriven adversarial learning for fundus image segmentation. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 102–110. Springer (2019)
22. Wang, W., Chen, C., Ding, M., Yu, H., Zha, S., Li, J.: Transbts: Multimodal brain tumor segmentation using transformer. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 109–119. Springer (2021)
23. Warfield, S.K., Zou, K.H., Wells, W.M.: Simultaneous truth and performance level estimation (staple): an algorithm for the validation of image segmentation. IEEE transactions on medical imaging 23(7), 903–921 (2004)
24. Whang, J., Delbracio, M., Talebi, H., Saharia, C., Dimakis, A.G., Milanfar, P.: Deblurring via stochastic refinement. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 16293–16303 (2022)
25. Wolleb, J., Sandkühler, R., Bieder, F., Valmaggia, P., Cattin, P.C.: Diffusion models for implicit image segmentation ensembles. arXiv preprint arXiv:2112.03145 (2021)
26. Yu, S., Xiao, D., Frost, S., Kanagasingam, Y.: Robust optic disc and cup segmentation with deep learning for glaucoma detection. Computerized Medical Imaging and Graphics 74, 61–71 (2019)
27. Zhao, R., Shi, Z.: Text-to-remote-sensing-image generation with structured generative adversarial networks. IEEE Geoscience and Remote Sensing Letters 19, 1–5 (2021)
28. Zhao, R., Shi, Z., Zou, Z.: High-resolution remote sensing image captioning based on structured attention. IEEE Transactions on Geoscience and Remote Sensing 60, 1–14 (2021)