VerseDiff-UNet

在医学图像分割扩散模型中引入形状先验模块

摘要

医学图像分割是诊断和治疗脊柱疾病的关键。然而，高噪声、模糊性和不确定性的存在使得这项任务极具挑战性。诸如不清楚的解剖边界、类间相似性和不合理的注释等因素导致了这一挑战。实现准确和多样化的分割模板是必不可少的，以支持放射科医生在临床实践。近年来，去噪扩散概率建模(DDPM)已成为计算机视觉领域的一个重要研究课题。它已经在各种视觉任务中证明了有效性，包括图像去模糊、超分辨率、异常检测，甚至是像素级的语义表示生成。尽管现有的扩散模型在视觉生成任务中具有鲁棒性，但它们仍然与离散掩模及其各种效果作斗争。为了满足对准确和多样化的脊柱医学图像分割模板的需求，我们提出了一个端到端的框架，称为VerseDiff-UNet，它利用了去噪扩散概率模型(DDPM)。我们的方法将扩散模型集成到一个标准的u形架构中。在每一步中，我们将添加了噪声的图像与标记的掩模相结合，以准确地引导扩散方向向目标区域。

此外，为了在医学图像中捕获特定的解剖学先验信息，我们结合了一个形状先验模块。该模块有效地从输入的脊柱图像中提取结构语义信息。我们在通过x射线成像获得的脊柱图像的单一数据集上评估我们的方法。我们的研究结果表明，VerseDiff-UNet在保留自然特征和解剖学变化的同时，在准确性方面显著优于其他最先进的方法。通过大量的实验验证了模型的通用性和有效性。通过对解剖结构进行更精确的分割，我们的框架有可能通过对解剖结构进行更精确的分割来促进对疾病的准确诊断和治疗。

2 方法

在本节中，我们将简要介绍我们提出的扩散建模框架。图1显示了我们提出的VerseDiff-UNet的训练阶段和采样阶段。与传统医学图像分割方法直接输入原始图像数据预测相应的分割标签图不同，扩散模型以原始图像和带噪声的分割标签图为输入学习去噪过程。

A .标签和时间嵌入

我们为转换后的多通道标签加入连续的阶跃噪声ε，称为扩散正演过程。

在得到带有t阶噪声的标签地图后，我们的目标是通过去噪模块基于标签地图和原始图像数据预测清晰的标签地图。

B.去噪模块

图1所示。提出的整体网络架构示意图

如图1所示，我们提出的VerseDiffUNet中的UNET也包含两个部分，一个编码器和一个解码器。首先，给定原始图像数据，其中N为模式图像的数目，将I和有噪声的单热标签xt，通道级联到VerseDiff-UNET的编码器中以获得多尺度特征IE。输入到网络的信息不仅通过传统的下采样和上采样结构提取，还通过逐层形状先验知识从Sf模块中提取信息，并与同级对应的解码特征ID相加求和，得到多尺度融合特征If。由于相同级别的编码器特征IE和解码器特征ID包含相同数目的相同大小的特征，因此可以通过与全局形状先验交互来进一步增强跳过的原始特征，从而便于生成具有区别性形状表示和全局上下文的特征，从而帮助模型更好地提取原始医学图像的脊椎解剖先验特征，最终导致预测结果x0。对于医学图像分割任务，我们的网络直接生成预测模板作为输出。其中，阶跃索引t与增加的嵌入和解码特征相结合。在其中的每一个中，它都使用共享的学习查找表来嵌入。

C.训练和抽样架构

如图2所示，VerseDiffUNet的主要架构是基于UNT的编解码器。编解码器由三个相应的下采样和上采样模块组成。扩散模型使用时间步长对马尔可夫过程执行变分推理，以学习训练数据分布p(x0)。在训练阶段，对于每一步，通过将输入图像添加到噪声分割掩模中来归纳解剖结构。具体地说，训练阶段输入具有手动标记掩模的原始脊椎图像的模型对，并且逐渐增加噪声。该框架包括一个正向过程和一个反向过程，也称为训练和采样过程。在对T中的每个时间步长的前向过程中，将高斯噪声添加到图像∼p()，直到图像变为各向同性的高斯。这种前向噪声过程表示为

其中(x0,x1,x2,.....xT)表示马尔可夫链中的时间步，α是控制噪声的方差的噪声调度器。在逆过程中，利用具有UNET网络结构的马尔科夫链构造增量去噪序列，以获得干净的图像。用不同的高斯噪声采样n次，产生n个不同的似然掩码。参数：

这是通过最小化所有时间步长的正向分布和反向分布之间的KL发散而获得的。

图2.a)采样和b)训练过程的图形模型

损失可以表示为：

迭代次数从均匀分布中采样，𝐠从高斯分布中采样。