CTS: A Consistency-Based Medical Image Segmentation Model
摘要:
在医学图像分割任务中,扩散模型已显示出巨大的潜力。然而,主流的扩散模型存在采样次数多、预测结果慢等缺点。最近,作为独立生成网络的一致性模型解决了这一问题。与扩散模型相比,一致性模型可以将采样次数减少到一次,不仅能达到类似的生成效果,还能大大加快训练和预测速度。然而,它们并不适合图像分割任务,在医学影像领域的应用也尚未得到探索。因此,本文将一致性模型应用于医学图像分割任务,设计了多尺度特征信号监督模式和损失函数引导,以实现模型收敛。实验验证了 CTS 模型在测试阶段只需一次采样就能获得较好的医学图像分割结果。
1 Introduction
医学图像分割领域一直是图像分割领域的热门研究方向。与传统的分割方法不同,利用生成模型进行图像分割也能取得良好的效果。由于扩散模型是一种从高斯噪声中采样的生成模型,其生成的图像具有很强的抗噪性和平滑性。因此,越来越多的研究利用扩散模型来解决不同图像的非生成问题。研究人员使用遮罩作为生成模型采样的目标,同时在生成模型中加入约束条件,以引导模型生成的方向。然而,由于在训练和预测过程中需要大量的重新采样,扩散模型计算效率低的问题亟待解决。 此外,在减少采样次数的同时,一致性模型还能确保采样的有效性。与 DDPMs相比,一致性模型代表了一种更优越的生成范式,但目前将该模型应用于医学图像分割领域的研究还很缺乏。因此,本文提出基于一致性模型构建医学图像分割模型,并根据分割损失和一致性训练损失设计损失函数,实现模型的端到端训练。
-我们构建了一个基于一致性模型的医学图像分割模型,其特点是采用了新设计的联合损失函数。
-在解码阶段,利用多尺度特征监督信号来引导模型的收敛方向。
2 Related Works
在本节中,我们将简要介绍与我们的工作相关的现有研究方向。扩散模型已被应用于许多领域,如序列建模[12,5]、语音处理[14]、计算机视觉[16,9]、计算机断层扫描(CT)和磁共振成像(MRI)等。在计算机视觉领域,为了减少采样次数,许多方法都做出了巨大努力。此外,还有一些专门针对条件生成的采样算法、如无分类器引导[11] 或有分类器引导[6]。图像分割是计算机视觉中的一项重要任务,它通过将图像分解为多个有意义的图像片段来简化图像的复杂性[2,8]。由于所需的时间、成本和专业知识[1,4],医学图像分割的图像和标签数量有限。因此,扩散模型通过综合标记数据,无需像素级标记数据,已成为图像分割研究中一种很有前途的方法。BrainSPADE[7]提出了一种合成带标记脑磁共振图像的生成模型,可用于训练分割模型。然而,医学图像分割中的扩散模型面临着采样次数多、预测时间长等问题。
3 Method
本文旨在充分发挥一致性模型一次采样的优势,同时保留分割模型的优点。在一致性模型中,直接训练一致性模型的方法被称为一致性训练损失,这也是 "CTS "名称中 "CT "的由来。具体过程如图 1 所示。
与一致性模型类似,本文的基本框架包括两个部分:模型 M 和目标模型 T M。模型的采样从每幅图像的掩码 xm 开始,输入相应的数据 xd 作为监督信号。初始化两个模型的参数,并将参数从 M 复制到 T M。
多尺度特征监督信号。多尺度特征监督信号 S xd i 的整合过程如图 1(a)所示。图像数据编码网络的解码器阶段逐步生成各种尺寸的特征图,并将它们与相应的监督信号相结合。在图像数据编码网络的解码器阶段,每个尺寸的特征图都会逐步生成相应的监督信号 xd i。监督信号 xd i 如图 1(b)所示。这一过程通过信道注意机制集成到 M 模型中,实现了多尺度监督信号的添加。在图像数据编码网络的解码器阶段,这些特征图包含不同尺度的信息,可以帮助模型更好地理解图像的细节和上下文信息。为了更好地整合监督信号和特征图,我们采用了通道关注机制。它可以自动学习每个通道的重要性权重,从而更好地利用监督信号中的信息。