当前位置：首页 > news >正文

SUnet: A multi-organ segmentation network based on multiple attention【医学图像分割】

news 2025/6/22 22:20:37

一、论文信息

1.1、中文名称

名称：SUnet：基于多重注意力的多器官分割网络

1.2、论文关键词

医学图像分割、Transformer、注意力机制、高效特征融合模块

1.3、核心概述

本文提出了一种新颖有效的医学图像分割方法 SUnet，用于腹部和胸部的多器官分割。

二、摘要

2.1、背景

腹部或胸部计算机断层扫描 (CT) 图像中的器官分割在医学诊断中发挥着至关重要的作用，因为它使医生能够快速定位和评估器官异常，从而指导手术计划并辅助治疗决策。

2.2、提出新方法

SUnet 是一个完全基于注意力的神经网络。首先，引入高效的空间缩减注意力（ESRA）模块，不仅可以更好地提取图像特征，还可以减少整体模型参数，并减轻过度拟合。其次，SUnet基于多重注意力的特征融合模块可以实现有效的跨尺度特征融合。此外，还考虑使用分组卷积和残差连接来增强注意力门（EAG）模块，提供更丰富的语义特征。

2.3、结果

本文评估了所提出的模型在突触多器官分割数据集和自动心脏诊断挑战数据集上的性能。SUnet 在这两个数据集上的平均 Dice 分别为 84.29% 和 92.25%，优于具有相似复杂性和规模的其他模型，并取得了最先进的结果。

三、引言

3.1、引出背景

腹部和胸部是人体脏器最多的部位，也是各种疾病的易发部位。从计算机断层扫描 (CT) 扫描中分割器官在诊断和治疗中发挥着至关重要的作用。然而，这对于医生来说是一项费力且容易出错的任务。因此，临床实践中迫切需要自动化的器官分割方法来辅助医生更高效、准确的诊断。近年来，人工智能虽然在生物医学领域取得了可喜的成果，如单细胞多组学数据分析、RNA-RNA相互作用、蛋白质组学研究、生物标志物研究等。

3.2、引出挑战

研究发现基因/蛋白质信号网络和药物代谢组学数据处理，腹部或胸部器官的自动分割仍然是一项具有挑战性的任务。 造成这些挑战的因素有很多，包括周围组织的干扰、器官变形或移位，以及导致边界不清晰的低图像对比度。这些挑战给实现准确和稳健的器官分割带来了困难。

3.3、目前研究概况

此前，已经提出了几种基于卷积神经网络（CNN）的二维（2D）医学图像分割模型，其中Unet模型是最具代表性的模型。 Unet 采用独特的 U 形编码器-解码器结构和跳跃连接，从而提高了模型的性能，同时引入了新的医学图像处理设计方法。 Unet 的出色性能随后激发了多种变体网络的发展。例如，Unet++ 通过使用嵌套和密集的跳跃连接代替原始连接来增强模型性能。 Unet3+是一种U形医学图像分割模型，采用全尺寸跳跃连接和深度监督。此外，ResUnet和ResUnet++也是基于Unet的优秀医学图像分割模型。然而，尽管取得了成就，这些基于 CNN 的模型仍遇到了问题，由于卷积运算固有的限制，有效建立长距离依赖关系的局限性，限制了其性能。

Transformer以其有效捕获长距离依赖的能力而闻名，引起了计算机视觉研究人员的关注。 Vision Transformer（ViT）是第一个将Transformer纳入图像处理的深度学习模型，它将图像划分为小块，并采用自注意力进行特征提取。 ViT 中的Transformer表现出了卓越的性能，为研究人员提供了一种新的特征提取方法，克服了卷积的固有限制。 TransUnet 率先将 Transformer 集成到医学图像分割中，用 CNN-Transformer 混合模型替代了 Unet 模型的编码器。这种混合特征提取网络可以提取更多语义特征，从而提高模型的性能。然后，Cao等人提出的Swin-Unet模型是一种完全基于Transformer的编码器-解码器架构，用于医学图像分割，已经取得了显着的性能。同样，几种基于Transformer的医学图像分割模型已应用于各种形式的医学数据。然而，这些基于 Transformer 的模型通常比基于 CNN 的医学图像分割模型具有更重要的参数和计算要求，这对计算资源有限的从业者构成了巨大的挑战。受pyramid vision transformer（视觉金字塔Transformer，PVT）的启发，本文设计了一种新颖的编码器-解码器医学图像分割网络，称为SUnet。该网络使用高效的空间缩减注意力（ESRA），从而提高性能并减少模型参数。此外，本文提出了一种高效的多注意力特征融合模块，该模块可以有效地将来自跳过连接的低级语义特征与来自解码器的上采样得到高级语义特征融合在一起。

3.4、贡献总结

本文提出SUnet，一种纯基于Transformer的U形医学图像分割网络，结合了高效的空间缩减注意力和多注意力特征融合。
本文提出了有效的空间缩减注意力，这使得模型能够在保持更少参数的同时表现更好，并减轻基于Transformer的模型中常见的过度拟合。
为了降低计算复杂度和数据依赖性，并提取更多与任务相关的特征，本文提供了基于分组卷积和残差连接的增强型注意门（EAG）模块。
本文提出了一种基于多重注意力的高效特征融合（EFF）模块，该模块在U形网络中实现了跳跃连接和解码器特征之间更好的融合。

四、相关工作

4.1、U-Net

Unet是Ronneberger等人提出的基于深度学习模型的卷积神经网络。 2015年它广泛用于医学图像分割任务，例如腹部多器官分割、自动心脏诊断、视网膜血管分割和皮肤癌分割。 Unet模型的特点是编码器-解码器结构和跳跃连接，将编码器的多级语义特征传输到解码器。这种操作可以将更多低级语义特征集成到特征图中，确保医学图像分割的高精度。由于其在图像分割方面的强大性能，它也广泛应用于卫星图像分割和工业缺陷检测。随着时间的推移，Unet演变出了许多深度学习模型。 本文提出的SUnet医学图像分割模型也继承了Unet的U型结构和跳跃连接。

4.2、注意力机制

注意力机制是深度学习中广泛使用的一种算法，它模仿人类的注意力机制。在深度学习中，所有输入特征都被同等对待，无论它们与当前任务的相关性如何。因此，该模型无法有效地将注意力分配到任务相关区域，这在一定程度上限制了深度学习模型的性能。相比之下，注意力机制对输入的不同部分分配不同程度的注意力和权重，更加关注与当前任务相关的信息。 空间注意力（SA）机制是深度学习领域广泛使用的技术。 2018 年Oktay 等人提出了注意力门（AG），它可以关注计算机视觉任务中的通道信息，它利用学习权重系数对原始输入与所选向量的乘积进行加权，实现通道选择和加权。 Hu等人提出的挤压和激励神经网络（SENet）。主要学习卷积神经网络中通道之间的相关性，并将较大的权重分配给对当前任务有用的通道。它通过两个步骤实现：挤压和激励。 Wang等人提出了一种有效的通道注意（ECA）模块。 ECA网络通过对每个通道的特征图应用一维卷积运算来建立通道之间的局部上下文相关性，从而实现通道注意力的自适应计算。除了上述仅关注特定维度的注意力机制外，还探索了多维度注意力机制。例如，Wu等人提出了一种多维混合注意卷积块注意模块（CBAM），重点关注通道和空间信息。与专门考虑通道维度的 SENet 注意力机制不同，CBAM 整合了多个维度的信息，以更好地关注有用信息。此外，Rahman等人首次首次提出了一种层次化的级联基于注意力的解码器，这也为EFF模块的设计提供了思路。

4.2、视觉金字塔Transformer

PVT是Wang等人提出的骨干网络, PVT的核心模块涉及通过卷积对多头自注意力（MHSA）机制中的键和值进行特征压缩。这种压缩操作过程显着降低了参数和计算复杂性。缩减率决定了卷积核和步幅的大小。 PVT 作为骨干网络广泛应用于各种视觉任务，例如目标检测和定位、遥感图像分类和医学图像分析。 PVT 的广泛采用凸显了其多功能性和有效性，使其成为计算机视觉领域具有重要意义的研究课题。

五、具体方法

5.1、框架整体结构

图1：SUNet架构图

SUnet遵循Unet的编码器-解码器结构，并采用跳跃连接来传达低级语义信息。通道号表示为 $C_{1}$ 、 $C_{2}$ 、 $C_{3}$ 和 $C_{4}$ ，定义为 $C_{1}$ = 64、 $C_{2}$ = 128、 $C_{3}$ = 320、 $C_{4}$ = 512。ESRA Transformer块用作 SUnet 中的特征提取模块，本文将两个通道堆叠在一起每级都有 ESRA Transformer块。 ESRA Transformer块的示意图如图1（b）所示。 ESRA Transformer块内的 Mix-FFN 模块如图1（c）所示，它与传统的前馈网络（FFN）不同，它在两个线性层之间使用深度卷积。在 SUnet 模型中，重叠补丁嵌入使用内核大小为 7 × 7 且步幅为 3 的卷积层进行补丁嵌入。重叠补丁合并采用内核大小为 3 × 3 的卷积层。嵌入补丁的重叠属性有助于减轻传统补丁嵌入造成的信息丢失。图1（d）展示了EFF的架构，主要由三个子模块组成：EAG、ECA和SA。本文使用组数为 32 的分组卷积和残差连接来改进原始 AG。 EAG模块通过上采样获得的高级语义特征来增强通过跳跃连接传输的低级语义特征。 EAG的结构如图1(e)所示。连接后，ECA和SA主要用于突出特征图中任务相关区域的重要通道和空间位置，以提高特征表达能力。应该注意的是，底部ECA和SA模块仅包含单个输入特征，因此本文仅使用ECA和SA来强调特征。

5.2、ESRA Transformer模块

图2：self-attention和ESRA的结构对比图

由于其强大的全局建模能力，Transformer 在计算机视觉任务中获得了广泛的应用。然而，当使用有限的数据进行训练时，基于 Transformer 的模型通常会遇到计算复杂性高和容易过度拟合等挑战。为了应对这些挑战，本文提出了一种称为 ESRA 的方法，如图 2 所示。ESRA不仅缓解了模型过拟合，还减少了整体参数。具体来说，本文利用卷积运算来压缩 MHSA 中的键和值，从而减少模型参数。原始MHSA的参数个数可以用公式(1)表示。

其中H和W表示输入特征图的高度和宽度，head是MHSA中的头的数量， $dim_{head}$ 表示每个头的通道尺寸。这里本文默认batch size为1。ESRA的参数计算遵循相同的原理，可以用公式（2）表示。

式中， $R_{i}$ 表示第i阶段的减速比。从式（1）和式（2）可以看出，当 $R_{i}$ > 2时，MHSA的参数数量比ESRA多。

为了减轻由 Transformer 引起的过拟合问题，本文在 ESRA 中使用 dropkey 来隐式为每个注意力头分配一个自适应算子。这种方法有助于通过惩罚具有较高注意力值的区域来限制注意力分布，促进更平滑的注意力并鼓励模型关注与任务相关的其他地方，捕获强大的全局特征。因此，ESRA可以表示为：

SR()运算可表示如下：

在等式(4)中，Conv2dRi表示通过使用核大小 $R_{i}$ 和步幅 $R_{i}$ 的2D卷积运算实现的特征压缩。而dropkey操作可以表示如下：

其中x表示要处理的注意力权重。函数bernoulli() 用于生成遵循伯努利分布的样本，而like() 函数则生成与x 大小相同的矩阵。

5.3、高效特征融合模块

AG模块是由Oktay等人提出的。在注意力网络中。然而，本文发现AG模块不仅计算复杂度高，而且需要很强的数据依赖性。当AG应用于高分辨率图像时，会显着增加计算负担。此外，AG 的两个输入之间必须存在严格的数据依赖，才能准确捕获重要特征。当输入之间的相关性较弱时，AG 无法捕获关键特征。在本文的工作中，通过用分组卷积代替传统卷积来扩展AG，以进行组内特征融合，并且分组卷积的计算复杂度明显小于传统卷积方法。同时，本文修改了AG的结构，在输入特征的卷积后添加ReLU层，并对跳跃连接传递的低级语义特征进行残差连接。当两个输入特征之间的相关性较弱时，残余连接可以减轻高级语义特征对低级语义特征的影响，从而避免整体模型的性能下降。 EAG的内部结构如图1（e）所示。 EAG可表示为：

其中Sigmoid和ReLU是激活函数，BN是Batch Normalization操作，GroupConv32是32组的分组卷积，Conv1×1是卷积核大小为1×1的常规卷积。在该模型中，g是通过以下方式获得的语义特征上采样，x 是跳过连接传递的低级语义特征。

5.3.1、基于多注意力的EFF模块

EFF模块的结构如图1(d)所示。 在EFF模块中，首先通过EAG增强两个不同级别的语义特征，以削弱不相关区域的影响。级联后通道数是原来的2倍，如果直接操作可能会丢失大量图像信息，本文使用ECA和SA从两个维度强调相关特征。值得注意的是，在本文的模型中，ECA 和 SA 串联连接。这种通道注意力和空间注意力的结合可以更好地实现多注意力融合。 ECA和SA结合使用，结构图如图3所示。

图3：ECA和SA串联结构

六、实验过程

6.1、数据集

Synapse 多器官分割数据集 (Synapse)：在本实验中，本文使用了 MICCAI 2015 多图集腹部 Labeling 挑战赛中的 30 张标记腹部 CT 扫描和 3779 张增强腹部图像。每次 CT 扫描由 85-198 个切片组成，分辨率为 512 x 512 像素。本文对八个不同的器官进行了图像分割：主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾脏和胃。

自动心脏诊断挑战数据集 (ACDC)：该数据集广泛用于心脏磁共振成像 (MRI) 评估，提供全面且完整注释的心脏 MRI 扫描集合。该数据集包含 100 名不同患者心脏的 MRI 扫描，每个样本包含三个器官标签，称为左心室 (LV)、右心室 (RV) 和心肌 (Myo)。在图4中，显示了部分数据集图像。

图4：部分数据集图像

6.2、实验细节

本文所有实验均基于Pytorch 1.8.0框架。使用Ubuntu 18.04操作系统、CPU I7-12700K、Nvidia RTX 3090、1 TB固态硬盘的计算机进行实验。在SUnet的所有实验中，本文都使用AdamW优化器，学习率和权重衰减设置为1e-4。

为了进行比较，本文在模型和所有基准模型中使用了相同的超参数设置。在Synapse数据集的实验中，数据被分为由18个样本数据组成的训练集和由12个样本数据组成的测试集。本文将批量大小设置为 24，最大 epoch 数设置为 150，输入图像大小和块大小分别设置为 224 × 224 和 16。应用随机翻转和旋转来增强数据。

在ACDC数据集上的实验中，本文使用70个扫描样本进行训练，10个扫描样本进行验证，20个样本进行测试。本文将批量大小设置为 12，epochs设置为 150，补丁大小设置为 16。本文使用随机翻转和旋转来增强数据。使用Dice损失和交叉熵损失函数，模型的整体损失可以表示为：

其中， $\lambda_1$ = 0.6， $\lambda_2$ = 0.4，DICE表示骰子损失函数，CE为交叉熵损失函数。

6.3、实验结果

6.3.1、Synapse数据集实验结果

表1：不同模型在Synapse数据集上的结果

图5：9种语义分割模型在Synapse上的性能

图6：SUnet 模型在 Synapse 数据集上的结果

图7：Synapse数据集上五个模型的定性实验

6.3.2、ACDC 数据集实验结果

表2：ACDC数据集中不同模型的比较

图8：SUnet 模型在 ACDC 数据集上的结果

图9：SUnet模型在ACDC数据集上的可视化结果

6.4、消融实验

表3：Synapse数据集上消融实验

七、总结

本文介绍了 SUnet，一种基于 ESRA 的新型 2D 医学图像分割模型。论文提出了一种创新的 EFF 模块，该模块使用多种注意机制（包括 EAG、ECA 和 SA）有效地融合跳跃连接和解码器特征。 EAG模块基于分组卷积，能够实现高效的组内特征融合。与其他 2D 医学图像分割模型（例如 TransUnet 和 Swin Unet）相比，本文提出的 SUnet 模型以更少的参数实现了更高的精度。它在 Synapse 数据集上的平均 Dice 为 84.29%，在 ACDC 数据集上的平均 Dice 为 92.25%。 SUnet 比当前最先进的 2D 医学图像分割模型表现出卓越的适应性和参数效率，使其更适合 2D 医学图像分割中的各种任务，特别是在计算资源有限的场景中。

尽管使用了多种注意力机制来实现特征融合，但SUnet并没有从根本上有效解决局部和全局特征融合的问题。在未来的研究中，将开发一种新的高效语义分割模型，该模型集成全局和局部图像特征，以提取更有效的图像特征，以改进医学图像分割。例如，未来的目标是通过使用CNN和Transformer形成一个新的混合特征提取单元，它对全局和局部特征同样关注。挑战在于实现这种更高效的混合模型，同时确保其不会显着增加参数甚至变得更加轻量级。通过解决这些挑战，渴望提高医学图像分割模型的性能，并为计算机辅助诊断和治疗做出贡献。