SwinEA：融合边缘感知的医学图像分割网络

叶晋豫，李 娇，邓红霞，张瑞欣，李海芳

(太原理工大学 信息与计算机学院，山西 晋中 030600)

0 引 言

关于医学图像分割[1，2]的研究方法层出不穷，最先进的医学图像分割框架大多基于U-Net[3]或其变体[4]。遵循这一技术路线有许多算法，如U-Net的各种变体[5-8]，通过跳跃连接保存了编码器层的单粒度信息，但忽略了丰富的多尺度空间信息，使得大量边缘信息丢失，这极大地影响了其在分割任务中的性能。

最近，受Transformer在自然语言处理方面取得巨大成功的启发[9]，其也应用于图像处理领域[10，11]。利用Transformer的关注全局的关键优势结合CNN的局部特性来构建一个更优的分割网络。目前，初步研究已经尝试将Transformer在医学图像分割领域应用[12-14]。一个常见的策略是将视觉Transformer直接应用于医学图像。ViT[15](vision Transformer)是第一个纯粹基于Transformer的架构。Touvron等[16]描述了几种培训策略，使ViT在ImageNet数据集上训练良好。Transformer也被应用于医学图像分割领域[17-22]，利用了Transformer和CNN的特性来提高模型的分割能力。例如，Chen等[2]和Hata等[1]将二维或三维解码器添加到类似ViT的编码器中。Hu等[11]使用类似Swin Transfomer[23]的编码器-解码器结构用于分割。现在，CNN和Transformer架构都展示了各自的独特优势，结合两种架构的优势进行综合分析是非常有意义的。

所以，本文提出了一种混合模型SwinEA(Swin Transformer fusion edge-aware medical image segmentation network)，即上下文金字塔结构结合Swin Transformer的网络，基于Swin Transformer的边缘感知上下文金字塔融合多尺度特征的医学图像分割网络。该网络利用上下文金字塔结构融合捕捉丰富的多尺度上下文信息，计算相邻像素间的局部相关性的卷积，用于提取边缘和角落等局部特征；并且结合设计用于序列到序列预测的Swin Transformer，利用其具有的先天性全局自注意力机制的替代体系结构，用来解决CNN的归纳偏置问题。

1 本文方法

本文提出的SwinEA的总体架构如图1(a)所示。SwinEA由编码器、解码器和边缘感知模块组成。其中编码器、解码器基本单元是Swin Transformer块。

图1 SwinEA模型框架

首先，编码器要将输入的医学图像分割为4×4的非重叠补丁块，再经过一个线性嵌入层(linear embedding)将特征维度投影到任意维度(表示为C)，转换后的补丁块通过4层Swin Transformer块和补丁合并层以生成分层不同尺度的特征。具体来说，补丁合并层(patch merging)负责下采样和增加维度，Swin Transformer块负责特征的学习。受U-Net[3]的启发，设计了一个基于Swin Transformer的对称解码器，该解码器由Swin Transformer块和补丁扩展层(patch expanding)构成，与补丁合并层相反，补丁扩展层专门设计用于执行上采样，补丁扩展层将相邻维度的特征图重塑为一个2倍分辨率的上采样特征图，再使用最后的补丁扩展层进行4倍上采样，恢复图像的分辨率并将特征映射到输入的分辨率大小(W×H)，然后再将这些上采样的特征输出像素级分割预测。

在编码器解码器结构的中间，本文不以跳跃连接进行连接，而是引入了一种上下文金字塔模块进行连接，将提取的上下文特征融合多尺度的特征通过边缘感知(edge-aware module)模块，以弥补损失多尺度空间信息，更好提取边缘和角落等局部特征，通过重构跳跃连接为解码器提供不同级别的全局上下文信息。

1.1 Swin Transformer模块

Swin Transformer块与传统的多头自注意(MSA)模块不同，它是基于滑动窗口构造的，由两个连续的Transformer组成，结构如图1(b)每一个Swin Transformer模块由LN(LayerNorm)层、多头自我注意模块、残差连接和具有GELU非线性的2层MLP(多层感知器)。两个连续的Transformer块分别采用基于窗口的多头自注意(W-MSA)模块和基于移动窗口的多头自注意(SW-MSA)模块。基于这样的窗口划分机制，Swin Transformer块可以表示为

(1)

(2)

(3)

(4)

self-Attention的计算方式如下

(5)

1.2 编码器

补丁合并层：输入的补丁分为4部分，并通过补丁合并层连接在一起。通过这样的处理，特征分辨率将被2倍下采样。由于连接操作导致特征维度增加了4倍，因此在连接的特征上应用线性层将特征维度统一为原始维度的2倍。补丁合并层无需卷积或插值操作即可实现下采样和特征维数的增加。

1.3 解码器

对应于编码器构建一个对称的解码器，它也是基于Swin Transformer块构建的。因此，对应于编码器中使用的补丁合并层，我们在解码器中构建了一个补丁扩展层，用于提取深度特征，进行上采样。补丁扩展层将相邻维度的特征图重塑为更高分辨率的特征图(2倍上采样)相应地将特征维度减少到原始维度的一半。补丁扩展层将相邻维度的特征地图重新塑造成一个2倍上采样分辨率的特征图。最后，利用网络结构最后一个补丁扩展层进行4倍上采样，恢复图像的分辨率到输入分辨率，然后经过线性投影层，将这些上采样特征输出像素级的分割预测。

1.4 边缘感知模块

为了解决固定尺度卷积操作在分割中的局限性，本文设计了边缘感知模块，通过不同尺度的空间信息提高了网络结构的鲁棒性。此外，U形网络中通常使用的跳跃连接，在接收字段不匹配会引入不相关的混乱和有语义差距。在本文中，提出全局上下文金字塔结构的边缘感知(EA)模块来解决这些问题，如图1(c)所示。在EA模块，会对跳跃连接进行重构。在解码过程中，浅层特征展示详细的边界信息，同时也会带来一些背景噪音。因此，边缘感知模块用于提取边缘特征，并进一步引导解码器，同时抑制浅层噪声并细化物体的轮廓。

在边缘感知模块中，利用不同大小的卷积核运算在多个分支中提取不同尺度的特征，通过融合不同粒度的空间信息，使网络能够学习更多的上下文信息。基于U形结构，在编码器和解码器之间设计了多个边缘感知(EA)模块，旨在通过重构跳转连接为解码器提供不同级别的全局上下文多尺度信息。众所周知，高层特征表达更多的语义信息，而浅层特征携带更多细节信息。因此，浅层深度主干的特征用于产生边缘特征，所以提出的边缘感知模块可以尽可能多提取浅层特征用于产生边缘特征。在EA模块中，将这一阶段的特征图与所有更高一级阶段的特征图相结合，重构跳过连接。

如图1(c)所示，以第二阶段的边缘感知模块为例，将各阶段的不同尺度的特征映射到与第二阶段相同的通道空间中的3×3卷积。接下来，生成的特征映射第三阶段特征地图和第四阶段特征地图上采样到与第二阶段特征地图相同的大小并连接。然后为了从中提取全局上下文信息不同层次的特征图，利用3种可分离卷积[24](Dsconv@1，Dsconv@2，Dsconv@4)，它们具有不同的膨胀率(1，2和4)进行平行使用，并且在这里采用可分离卷积可以减少模型参数。值得注意的是平行路径和扩张的数量速率随融合阶段的数目而变化。最后，通过卷积得到最终的特征地图。

以上，可以总结出各个阶段的EA模块(为了简化公式，忽略正则卷积)

(6)

其中，Gk为插入第k阶段，Fk为编码器中第k阶段的特征图，⊗2i-k为上采样操作速率为2i-k，c表示串接操作和Dsconv@2i-k为可分离的扩张卷积膨胀率为2i-k。

通过引入多个边缘感知模块在编码器和解码器之间，全局语义高层次的空间信息流可以逐层引导到不同阶段。

1.5 损失函数

医学图像分割的一个主要挑战是分类分布不平衡。为了进一步优化我们的模型，也为了更好地评估标签轮廓和预测的轮廓之间的相似性，我们将组合一个新的损失函数，由交叉熵损失和Dice Loss组成。

交叉熵损失函数lce计算如下

(7)

其中，σθ(x) 是深度神经网络的输出，x是网络的输入，而y代表标签。θ表示神经网络的可训练参数。

Dice Loss函数lDice计算如下

(8)

其中，|X∩Y| 表示X和Y集合的交集，|X| 表示元素个数，|Y| 表示分割任务中的标签及预测结果。

我们采用由Dice损失lDice和交叉熵损失函数lce计成的联合损失来执行所有分割任务。

公式如下

ltotal=λlDice+(1-λ)lce

(9)

其中，λ是Dice损失和交叉熵之间的权衡损失。为了公平比较，我们实验中的每个单独的任务都使用相同的损失函数来执行。

2 实验和结果分析

2.1 实验数据集

Synapse多器官分割数据集：该数据集包括3779张轴向腹部临床CT图像。训练集包括18个样本，测试集包括12个样本。其中每个CT图像样本由85～198个512×512像素的切片组成，体素空间分辨率为 ([0.54～0.54]×[0.98～0.98]×[2.5～5.0]) mm3。 本文对8个腹部脏器(主动脉、胆囊、左肾、右肾、肝、胰、脾、胃)进行实验验证得到其平均DSC和平均Hausdorff距离(HD)。

ACDC(心脏自动诊断挑战)：ACDC是一个公共心脏MRI数据集，包括100例样本。一系列短轴切片，覆盖心脏从左心室基部到心尖，厚度为5～8 mm。短轴平面内空间分辨率从0.83到1.75 mm2/pixel，对应的标签包括左心室(LV)、右心室(RV)和心肌(MYO)。数据集分为70个训练样本(1930个轴向切片)、10个验证样本和20个测试样本。

2.2 实验设置

2.2.1 环境

SwinEA网络是基于Python 3.6和Pytorch 1.7.0实现的。对于所有的训练案例，都使用数据增强，如翻转和旋转，以增加数据的多样性。输入的图像大小设置为224×224，patch大小设置为4。本文在内存为32 GB的Nvidia V100 GPU、Linux操作系统上训练本文的模型。

在训练期间，批量大小为24，使用的是SGD优化器，动量为0.9，权值衰减为1e-4，对本文的反向传播模型进行优化。实验设置300个epoch。

2.2.2 评估指标

本文使用两种类型的指标来评估所有模型。

评价预测与分割标签重叠程度的Dice相似性系数(DSC)

(10)

其中，P为预测的分割图，G为标签。

HD(Hausdorff distance)，测量两个分割贴图之间最大的对称距离

(11)

式中：d(·) 为欧几里得距离，sup和inf分别为上极值和下极值。P为预测的分割图，G为标签。本文采用95%HD，以消除一个非常小的子集的异常值的影响。

2.3 实验结果

2.3.1 在Synapse多器官CT数据集上的验证

在Synapse多器官CT数据集上，本文提出的SwinEA与之前较新的方法的比较见表1，其中2-3列分别为平均DSC和HD评价指标结果，后8列为8个器官的具体DSC指标结果。实验结果表明，本文的基于Swin Transformer融合多尺度上下文金字塔的方法分割精度达到了79.65%(DSC↑)和20.97%(HD↓)。与最近的TransUnet、Swin-Unet、MTUnet[25]方法相比，本文的算法在DSC评价指标和HD评价指标上有着提升，这表明本文的方法可以实现更好的分割边缘预测。

表1 不同算法在Synapse多器官分割数据集的评价指标对比结果

不同方法对Synapse多器官CT数据集的分割结果如图2所示。从图中白色框标出的地方可以看出，本文方法在一些分割边缘上比其它一些先进方法均有着分割边缘清晰，区域明显的效果。因此，在本文工作中，验证了基于Swin Transformer的上下文金字塔融合多尺度特征结合的医学图像分割网络通过边缘融合模块捕捉了丰富的多尺度上下文信息，计算相邻像素间的局部相关性的卷积，在提取边缘和角落等局部特征方面表现得很好，也获得更好的分割结果。

图2 不同算法在Synapse数据集上的分割对比结果

表1对不同算法的实验结果进行了定量评估。U-Net为原始深度卷积网络的方法；Transunet为生成模型的编码结构使用Transformer，解码器中使用卷积结构；SwinUnet为编解码器结构均使用纯的Transformer方法；SwinEA为本文提出的方法。从表中可以看出，通过使用上下文金字塔结构与Swin Transformer方法共同结合，DSC提升到79.65%，HD提升到20.97%。数据方面也验证，在Swin Transformer网络中引入上下文金字塔机制，可以有效且精确地针对边缘完成腹部多器官分割的任务。

为了验证该模型的性能，与分割标签通过三维重建进行了比较。将预处理的切片结果还原为nii格式的三维数据集，比较三维下的分割结果。

图3为模型在Synapse数据集不同层下的分割结果。第一列为Synapse数据集切片后的原始图像，第二列为标签结果，第三列为本文提出的网络模型的分割结果，第四列为本文进行三维重建的分割结果。从图3中可以看出，与第二列专家分割的结果进行对比，第三列本文的实验分割结果在细节上分割精确，基本上与标签分割结果相似。并且将切片数据进行三维重建之后，边缘圆滑。

图3 三维重建的Synapse分割结果

2.3.2 在ACDC数据集上的验证

为了验证模型的泛化性，将SwinEA在ACDC 数据集来进行医学图像分割。表2中总结了不同算法下的分割实验数据结果，其中第2列为平均DSC值，后3列为右心室、心肌、左心室的具体DSC值。使用MRI模式的图像数据作为输入，SwinEA仍然能够达到优秀的性能，平均准确率达到91.34%。这表明我们的方法具有良好的泛化能力和稳健性。如图4中，第一行为ACDC数据集切片后的原始图像，第二行为分割标签，第三行为本文提出的网络模型的分割结果，第四行为本文三维重建的分割结果。从图4中可以看出，通过与第二行标签结果的结果进行对比，第三行本文的实验分割结果在细节上分割精确，基本上与标签分割结果相似，并且将切片数据进行三维重建之后，可以看到分割边缘圆滑，分割区域清晰。

表2 不同算法在ACDC心脏分割数据集的评价指标结果

图4 不同算法在ACDC 数据集的分割结果

2.3.3 消融实验

为了探究不同因素对模型性能的影响，本文在Synapse数据集上进行了消融实验研究。

(1)边缘感知模块的位置的影响

本文的EA模块被添加到1/4、1/8和1/16分辨率刻度的位置在图中分别代表边缘感知模块1、边缘感知模块2、边缘感知模块3，探讨不同位置对分割效果的影响。从图5中可以看出模型的分割效果，边缘感知模块1、边缘感知模块2代表的黑色白色部分大部分表现比边缘感知模块3的灰色更准确。因此，为了使模型更加鲁棒，本文将不使用边缘感知模块3的情况。

图5 边缘感知模块的位置的影响

在实验过程中发现此模块在不同位置上的表现不同，位于浅层时效果更好，可以更好提取细节空间信息。在添加块数上，块数更多会随之带来提取信息过于冗余，表现一般，平衡位置与添加块数的平衡，在实验中验证添加两块EA模块，且位于浅层位置可以达到一个更优的结果。

(2)边缘感知模块的数量的影响

本文的边缘感知模块数量的讨论结合位置限制，所以最多可以有3块，最少可以为0块。探讨数量对分割性能的影响。从表3中可以数量看出模型的分割性能，3个边缘感知模块造成冗余，效果一般。两块效果普遍比一块效果更优。因此，为了使模型更加健壮，本文将采用两块边缘感知模块。

表3 边缘感知模块的数量的影响

本文通过消融实验，得出关于边缘感知模块模块数量影响的结论，添加3块边缘感知模块模块时造成提取信息过于冗余，造成效果不佳，而针对于分割边缘低维度信息特征更加有用，所以两块时，添加处于低维度特征的边缘感知模块1、边缘感知模块2模块时，效果更好，将高分辨率的低水平特征和低分辨率的高水平特征融合起来，帮助划分出详细清晰的分割边界。

(3)损失函数的影响

本文采用的损失函数包括常用在分割中的Dice损失和交叉熵损失，两项损失进行消融实验后，实验结果见表4，结果验证两项损失联合训练的方法可以得到更好的分割效果。针对分割边缘的清晰与精度，这两种损失提供了损失的信息，更好解决了固定尺度卷积操作在目标分割中的局限性。

表4 损失函数的影响

3 结束语

本文提出了一种针对分割边缘的利用Swin Transformer融合边缘感知的医学图像分割网络(SwinEA)，设计了一个基于上下文金字塔的边缘感知(Edge-Aware)模块，用于融合全局多尺度的上下文信息，主要针对边缘和角落等局部特征。并且侧重于解决U形网络中全局多尺度上下文信息捕获和整合的弱点。新型边缘感知模块以上下文金字塔为基础的边界感知模块插入U形框架，开发和融合丰富的全局多尺度上下文信息。将高分辨率的低水平特征和低水平的特征融合起来。高分辨率的低层次特征和低分辨率的高层次特征来帮助划定详细的分割边缘。

本文对不同类型的医学图像分割任务进行了综合实验，以验证所提出的SwinEA的有效性和泛化性，包括腹部多器官分割任务和ACDC数据集。本文方法可以通过进一步的处理达到更好的性能，未来可以扩展应用到3D图像以及其它医学图像分割任务，这也是我们近期的工作。