MedSegDiff: Medical Image Segmentation withDiffusion Probabilistic Model

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

摘要：

扩散概率模型(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]。

最近，扩散概率模型(diffusion probistic model, DPM)[9]作为一类强大的生成模型[27]得到了广泛的应用，它能够生成具有高多样性和高合成质量的图像。最近的大型扩散模型，如DALLE2[17]、Imagen[19]和Stable diffusion[18]已经显示出令人难以置信的生成能力。扩散模型最初应用于不存在绝对真值的领域。然而，最近的研究表明，它也有效地解决了地面真相是唯一的问题，如超分辨率[20]和去模糊[24]。

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

--首先提出了基于dpm的一般医学图像分割模型。

--提出了分步关注的动态条件编码策略。

--为了消除高频分量的负面影响，提出了高频解析器。

--SOTA在三种不同图像模态的医学分割任务中的性能。

2  方法

我们基于文献[9]中的扩散模型来设计模型。扩散模型是由正向扩散和反向扩散两个阶段组成的生成模型。在正演过程中，分割标签x0通过一系列步骤T逐渐加入高斯噪声。在反向过程中，通过对噪声过程进行反向，训练神经网络来恢复原始数据，可以表示为:

式中θ为反工艺参数。从高斯噪声出发，pθ(xT) = N (xT;0, In×n)，其中I为原始图像，反向过程将潜在变量分布pθ(xT)转换为数据分布pθ(x0)。为了与正演过程保持对称，反向过程逐步恢复噪声图像，以获得最终清晰的分割。

遵循DPM的标准实现，我们采用UNet作为学习网络。

图1:MedSegDiff的图解。为清晰起见，图中省略了时间步长编码。

图1显示了一个示例。为了实现分割，我们对步长估计函数xt通过原始图像先验，可以表示为:

其中EIt为条件特征嵌入，在本例中为原始图像嵌入，Ext为当前步骤的分割映射特征嵌入。将这两个分量相加并发送到UNet解码器D进行重构。步骤索引t与添加的嵌入和解码器功能集成在一起。

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

2.1  动态条件编码

在大多数条件DPM中，条件先验将是一个唯一的给定信息。然而，医学图像分割是出了名的模糊对象。病变或组织通常很难从其背景中区分出来。

低对比度的图像模式，如核磁共振成像或超声图像，使其更糟。只给定静态图像I作为每一步的条件将很难学习。为了解决这个问题，我们提出了每个步骤的动态条件编码。我们注意到，一方面原始图像包含准确的分割目标信息，但难以与背景区分，另一方面，当前步分割图包含增强的目标区域，但不准确。这促使我们去整合

当前步分割信息xt转化为有条件的原始图像编码进行相互补全。具体地说，我们在功能级别上实现集成。在原始图像编码器中，我们利用当前步长编码特征来增强其中间特征。条件特征图mkI的每个尺度与形状相同的xt编码特征mkx融合，k为层的指数。这种融合是通过一种类似于注意力的机制a来实现的。特别是，首先将两个特征映射应用层归一化并相乘以获得亲和映射。然后将亲和映射与条件编码特征相乘，增强关注区域，即:

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

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

2.2  FF-Parser

我们以特征集成的路径方式连接ff解析器。它的功能是约束xt特性中与噪声相关的组件。我们的主要思想是学习一个参数化的关注(权重)映射应用于傅里叶空间特征。给定一个解码器特征映射m∈RH×W ×C，我们首先沿着空间维度执行二维FFT(快速傅立叶变换)，我们可以表示为:

式中F[·]为二维FFT。然后，我们通过将一个参数化的关注映射a∈CH×W ×C乘以m来调制m的谱:m 0 = a⊗m，(5)，其中⊗表示元素积。最后，我们采用逆FFT将m0逆回空间域:m0 = F−1[m0]。(6) FF-Parser可以看作是频率滤波器的一种可学习版本，频率滤波器广泛应用于数字图像处理[16]。与空间关注不同，它对特定频率的分量进行全局调整。从而可以学会约束高频分量进行自适应积分。

图2:FF-Parser的示例。FFT表示快速傅里叶变换。

2.3  训练与结构

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

在每次迭代中，将随机抽取一对原始图像Ii和分割标签Si进行训练。迭代次数从均匀分布和∊来自高斯分布。

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

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

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

3 实验

3.1数据集

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

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

3.2实现细节

我们分别对MedSegDiff++、MedSegDiff- l、MedSegDiff- b和MedSegDiff- s模型的大型、大型、基本和小型变体进行了实验。

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

3.3 Main Results

我们将针对这三个具体任务提出的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)。我们可以看到，先进的网络架构通常会获得更好的结果。例如，在光学杯分割中，基于viti的通用分割方法:TransUNet在方法:BEAL上甚至优于基于cnn的任务。在脑肿瘤分割方面，最近提出的基于dpm的分割方法EnsemDiff优于之前基于vita的竞争对手TransBTS和TransUNet。MedSegDiff不仅采用了最近成功的DPM，而且针对一般的医学图像分割任务设计了相应的策略。我们可以看到MedSegDiff在三个不同的任务上都优于所有其他方法，这表明了MedSegDiff在不同的医学分割任务和不同的图像模态上的泛化。与专门用于脑肿瘤分割的基于dpm的模型(即EnsemDiff)相比，该模型在Dice上提高了2.3%，在IoU上提高了2.4%，这表明了我们独特的技术(即动态调节和FF-Parser)的有效性。

图3:表1中Top-4种常用医学图像分割方法的视觉对比。从上到下分别为脑肿瘤分割、光学杯分割和甲状腺结节分割。

在图像上，使它们难以被人眼识别。通过与这些计算机辅助方法的比较，可以明显地看出，该方法生成的分割图比其他方法更准确，特别是对于模糊区域。将DPM与所提出的动态调节和FF-Parser相结合，可以更好地定位和校准低对比度或模糊图像的分割。

表1:MedSegDiff与SOTA分割方法的比较。最好的结果用粗体表示。灰色背景表示针对该/这些特定任务提出的方法。

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解析器。

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