一种基于局部与全局表征交叉耦合的脑卒中病灶分割网络模型

2023-11-07倪波柯亨进刘志远

中南民族大学学报（自然科学版） 2023年6期

倪波，柯亨进，刘志远

（湖北理工学院计算机学院 &黄石市类脑计算与神经科学重点实验室，黄石 435003）

脑卒中（俗称脑中风）作为一种常见心脑血管疾病，致残率和死亡率高，预后症状包括意识丧失或言语模糊肢体功能障碍及丧失等［1］.脑卒中术前诊断及其术后评价主要依靠临床医生对患者脑部T1-Weighted MRI 扫描数据上的病灶位置、面积大小进行准确的定义.但患者发病程度不同，病灶区域在影像中呈现出的位置、大小也不一样等（如图1所示），导致许多分割方法对脑卒中病灶的分割效果还无法应用于临床实践.到目前为止，这项工作主要还是依靠有经验的临床医生在患者脑部T1-Weighted MRI扫描数据的每一帧数据上手动定义病灶区域轮廓，该方式即耗时也可能在定义过程中引入主观误差.因此，开发了一种精确的脑卒中MRI的分割方法对于制定脑卒中术前计划和术后康复方案具有重要临床价值.

图1 脑卒中T1-Weighted MRI病灶示例Fig.1 Example of stroke T1-weighted MRI lesion

当前卷积神经网络（Convolution Neural Network，CNN）是医学图像分割领域的主流方法，其中代表性工作有Unet［2］，Attention Unet［3］，Res-Unet［4］，Dense-Unet［5］.它们都是一种U形的编码器-解码器架构，其中编码器通过2D 卷积-池化操作自顶向下提取图像多尺度的层次化特征，然后使用一种跳跃连接将编码器提取的各层次特征输送至解码器进行逐层反卷积操作从而预测分割结果.大量实验证明这类架构能够有效增强图像在特征提取过程中由于卷积和池化操作所导致的目标细节的损失问题.文献［6-9］中运用3D 卷积操作对目标在CT/MRI 体数据中变化信息进行建模.此外，一些新的卷积变体，如空洞卷积［10］和变形卷积［11］也相继被提出用于上获取目标形变的局部特征.这些工作都可看作是提高CNN 获得更具判别性目标特征的一种尝试，但CNN对目标的感受野受限于卷积核大小，卷积核的有限尺度使其在构建目标与背景的长距离依赖关系依然具有局限性，而这种目标与背景之间的长距离依赖关系对高级语义理解任务至关重要.

Transformer［12］作为一种新型的自注意力机制被首次提出，该模型在构建特征的长距离依赖关系表现强大能力，并在机器翻译领域取得了比传统CNN模型更好的结果.此后，大量基于Transformer 的模型被应用于计算机视觉领域［13］，例如，目标检测［14］，语义分割［15］，视频分析［16］等.Transformer 的优势在于构建图像的全局上下文信息，但Transformer 是在具有固定大小的图像块之间进行逐像素计算，当面对大尺度图像的密集预测任务时整个训练或推理将会带来较大的计算成本而且具有固定大小的图像块对多尺度变化目标的上下文关系建模也会带来问题.而目标的局部信息和全局信息对实现医学图像精确分割精度具有很重要的意义.如何在全局信息和局部信息之间找到完美的平衡，是提升分割模型性能的重要问题之一.由于大多数医学图像数据集都很小，使得训练出兼顾两者的分割模型更具有挑战性.

当前将CNN 和Transformer 相结合成为一个热点.例如，MedT［17］，Cell-DETR［18］，Conformer［19］等.这些模型试图将基于CNN 的局部特征和基于Transformer的全局特征进行融合并用于图像分割任务.本文针对脑卒中病灶区域的特点，提出了一个新的结合Transformer 和CNN 的U 型网络.如图2 所示，本网络延续着基于编码器-解码器的U 型架构风格，其中编码器是基于Swin Transformer和CNN的双分支网络，该双分支网络分别提取输入图像的全局和局部特征.为了使这两个分支网络增强彼此的特征学习能力，每个分支将提取的特征表示输入另一个分支用于学习.此外，还设计了一种融合模块（Fusion Block）用于对编码器提取的全局特征和局部特征进行融合，以缓解随着网络的加深小目标、弱边缘等特征在层次化特征编码过程中被弱化的问题.融合后的特征表达被逐层反馈给一系列上采样器进行逐层解码，直至输出分割的预测结果.此外为了提高网络的收敛速度，我们提出了一个新的损失函数，它结合了Dice函数和Focal［20］损失函数的优点.这种新方法比使用传统的Dice 函数收敛得更快，并表现出更平滑的收敛曲线.总的来说，本文的主要工作可总结如下：

图2 本文网络的架构示意图Fig.2 Diagram of the proposed network

（1）提出了一种新型的结合全局特征和局部特征的U 型分割网络.该网络编码器是一个包含Swin Transformer 和一个CNN 的双分支网路.该双分支网络相互学习对方风格的特征表达，从而提取目标更具判别性特征信息.

（2）为了缓解特征提取过程中小目标以及弱边缘特征被弱化的现象，设计了一种融合模块对在编码器的每个阶段将局部特征与全局特征进行融合，然后将融合后的特征表达转发到一个级联的上采样器用于增强对病灶区域的预测能力.

（3）本文方法与相关医学图像分割方法在开源脑卒中数据集（ATLAS）［21］上开展了实验对比.结果显示，本文方法不仅在测试集上提供了最优的分割效果，而且针对小目标分割上也能够最大程度保留病灶区域的形状和边缘形象.

1 相关工作

1.1 基于U型架构的医学图像分割网络

自Unet［2］在分割领域取得成功以来，大量基于U型架构的模型被学者所提出，例如，Attention Unet［3］，Res-Unet［4］，Dense-Unet［5］，X-net［8］.这些方法主要是在网络中嵌入残差连接和密度连接使其网络在特征提取阶段能够更加关注待分割的目标以及在训练阶段缓解梯度消失的问题.Unet++［22］和Unet3++［23］则分别是在Res-Unet［4］和Dense-Unet［5］的基础上对跳跃连接的方式进行了优化从而缓解模型过拟合现象.文献［24］提出了一种卷积核注意力机制的模型用于提取多尺度特征的提取和通道融合，从而提高分割精度.文献［6，25-26］将3D 卷积操作应用于CT/MRI 体数据使得在获取二维目标特征的同时，也能够提取目标在三维空间中的变化信息，但3D卷积操作也会带来计算复杂性和过多的参数的问题，从而影响模型的训练和推理效率.Double-Unet［9］则是采取了2D卷积和3D卷积相结合的方式试图在获取更多目标特征和降低计算效率上寻求一种高效的平衡.总的来说，基于U 形架构的CNN 在医学图像上任然显示出一定发展潜力.

1.2 基于Transformer的医学图像分割网络

Transformer［12］作为一种新的自注意力模型，首次被提出并应用于机器翻译领域.ViT［27］是首个利用Transformer 应用于图像分类任务的工作，并且在多个数据集上的表现都不输最先进CNN 模型.随后，ROBIN 等［28］提出了Segmenter 模型，该模型是一种只使用Transformer 的语义分割方法.如ViT［27］一样，Segmenter 将图像被编码为一系列小的图像块.Trans Unet［29］将Transformer 作为与Unet 模型的编码器提取具有具有上下文关系的特征图，然后借助解码器对编码器输出特征进行上采样并与高分辨率的CNN 特征组合实现对分割目标的预测.实验证明以Transformer 模型在构建目标的全局特征方面能够提供比CNN 更好性能，但Transformer 的全局自注意力是基于逐像素计算，当面对高分辨率图像的语义分割任务时就会带来较高的计算代价.Swin Transformer［30］在此基础上的改进主要是将固定大小的图像块划分为几个同样大小的区域，然后只针对该区域的图像块实施自注意力机制，然后运用一种滑动窗口（Shift-Windows）的方式将自注意力计算限制在不重叠的局部窗口，同时还通过不同窗口的连接来获取图像的多层次特征表达.该方法在多个数据集上的实验证明表明，Swin Transformer 在构建多尺度的图像特征的同时也具有良好的计算效率，可以作为多类型视觉任务的骨干网络.随后，Swin-Unet［30］将Swin Transformer 引入U 型架构用于医学图像分割.

近来，将CNN和Transformer相结合成为医学图像分割领域的一个热点.例如，MedT［17］在Transformer［12］的基础上提出了一种轴向门控注意力模型，从而促使Transformer 能够在小数据集上获得更好的性能.Cell-DETR［18］则是在CNN 骨干模型上增加自注意力从而实现对细胞的实例分割.Conformer［19］中，学者构建了一个由CNN 和Transformer 组成的双分支网络，并提出了一种特征耦合单元用于融合基于CNN的局部特征和基于Transformer 的全局特征.还有一些基于Transformer［32-33］模型被提出用于腹腔器官的分割这些模型进一步显示了基于Transformer 自注意力在构建长距离依赖关系方面的能力，但也由于自注意力需要图像块中每个像素进行计算，从而导致计算量大的问题依然存在.

2 本文方法

2.1 网络架构

从图2 可看出，本网络与传统U-net架构不同的是编码器是一个双分支网络，其中一个分支是基于Swin Transformer 网络，另一个则是基于CNN 的网络.这两个分支网络分别负责提取输入图像的层次化全局特征和局部特征.与文献［30］类似，本文的Swin Transformer 分支网络也包含四个处理单元，每个单元都包含基于Swin Transformer 的自注意力模块（STrans Block）负责对输入数据实施基于Shift Window 的自制力机制（细节见第2.2 节）.对于第1个处理单元，输入图像X ∈RH×W×3首先会被“Patch&Linear Embedding”层划分为一系列不重叠的图像块Patch={xi∈RP2·C，i=1，…，N}，每个xi的分辨率是P×P，其像素值就是该图像块的特征.随后被转换到适合STrans Block 的输入维度C.STrans Block负责对这些计算Patch 之间的全局特征.为了构建与CNN 分支网络相匹配的层次化特征，后面处理单元中的图像块合并层（Patch Merging Layer）以2×2间隔对STrans Block 输出的特征图进行采样，输出维度被设置为2C.因此当输入图像通过第1、2、3、4个处理单元时，特征图的输出分辨率分别为

CNN 分支网络由一系列的卷积模块（Conv Block）组成，负责提取输入图像的层次化局部特征.该分支网络中的每个Conv Block 都采用了特征金字塔结构，特征图的分辨率随着CNN 网络深度的增加而降低，而通道数则增加.每个Conv Block 由两个BN-ReLu-Conv 层组成（Kernel Size=3，Stride=1，Padding=1）.在每个BN-ReLu-Conv 层之后，实施了一个Dropout 为0.3 的操作，以克服潜在的过拟合问题.需要说明的是，CNN 分支网络的第一个单元有两个Conv Block 主要是为了使该分支在该单元输出的特征大小与Swin Transformer 分支中第一个单元输出的特征图的大小相匹配.

随编码器加深，Swin Transformer 分支网络中合并层以及CNN 分支网络中的池化操作会使得图像中的小目标和模糊边缘特征会逐渐被弱化的现象，本文设计了一种融合模块（Fusion Block）对每个处理单元输出的局部特征和全局特征进行融合，并将融合后的特征图输送由一组级联的升采样器（Upsampler）组成的解码器，用于逐级恢复来自不同融合模块的特征图.如图3（c）所示，每个解码器块由多个上采样层组成，对来自Fusion Block的特征表示进行解码，以输出最终的分割掩码.具体来说，底部Upsampler 负责接收由编码器的第4 个Fusion Block 获得的融合特征图，它首先对特征图进行2×2的双线性上采样，使输入特征图的维度加倍.至于其余的解码器区块，我们首先将前一个上采样区块的特征图与相应的合并区块的输出连接起来.然后，对合并后的特征图进行同样的操作，如上采样、卷积和批量归一化与ReLU激活.

图3 部分模块的详细情况Fig.3 Details of some modules in the proposed network.

2.2 STrans模块

图3（a）显示了STrans Block 的结构，每个模块由一个基于窗口多头注意力模块（Window Multi-Head Self Attention，W_MHSA），一个基于移位窗口的多头自我注意模块（Shift Window Multi-Head Self Attention，SW_MHSA）、一个具有GELU 非线性的2层多层感知机（Multi-Layer Perception，MLP）和2 个Layer Norm Layer 组成.W_MHSA 和SW_MHSA 中的Shift Windows机制主要是为了对不同位置的图像块之间实施自注意力机制，从而构建图像块之间的多尺度自注意力特征.图4 说明了W_MHSA 和SW_MHSA 的工作机制，其中W_MHSA 表示从将STrans Block输出的第l个特征图从左上角开始划分为M×M 个不重叠的窗口区域（M=2），然后基于多头自注意力被应用在每个窗口区域.这样做的好处是自注意力的计算被限制在一个固定的窗口区域，提高了计算效率.在第l+1 层特征图中，窗口平移个位置，从而让不同窗口中的图像块相互计算.在STrans Block输出的第l和l+1层特征图的输出可以表示为：

图4 Swin Transformer中滑动窗口机制示意图Fig.4 Illustration of the Shift Window mechanism in Swin Transformer

其中Q，K，V∈RM2×d分别表示Query、Key 和Value三个矩阵，M2和d分别表示图像块的数量以及Query和Key 矩阵的维度.矩阵B中的值来自于偏置矩阵

2.3 Fusion Block

Fusion Block 主要负责对每个处理单元输出的局部特征图和全局特征图进行融合，从而弥合两种风格的特征图之间存在的语义鸿沟.图3（b）说明了每个Fusion 模块的结构首先通过一个Concatenation层对两个风格的特征图进行拼接，然后通过两个连续的BN-ReLu-Conv（3×3）层将它们融合.

2.4 损失函数

常见的语义分割损失函数是Dice 函数和二值交叉熵函数（Binary cross entropy，BCE）.但在本文任务中，对一些小出血点，病灶面积占背景比例较小，会导致传统损失函数在训练过程容易陷入局部最优.为提供更加精确的分割效果，我们结合Dice函数和Focal 函数［22］构建了一个新的损失函数来缓解前景-背景不平衡的问题，并将该损失函数表述为：

其中g∈0，1 分别表示图像中病灶区域和背景区域的标注，p∈[0，1]则表示分割算法计算的分割区域像素的概率值，其中“0”表示图像的背景值，“1”表示病灶区域的前景值.LFocal和LDice分别表示Focal Loss 函数和Dice 函数，LFocal是在二值交叉熵损失函数的基础上的改进，从而让损失值聚焦在分割区域，其表达式为：

其中Nf和Nb分别表示前景和背景的像素数量，α∈(0，1]和γ∈[0，5]分别是调节LFocal权重的系数.LDice方程表示是：

其中δ∈[0，1]是一个调节参数用于防止被零误除，并让负样本也有梯度传播.

3 实验

3.1 数据集和评价指标

选择一个ATLAS［21］作为验证本文方法性能的数据集.ATLAS［21］是一个有关脑卒中T1-Weighted MRI 开源数据集，包括955 个T1-Weighted MRI 脑部扫描数据和对应的病灶区域的Mask图，每个脑部扫描数据大小为233×197×189，其中233×197 是MRI 二维断层切片的大小，189 是每个脑部扫描的切片数量，物理像素大小为0.9×0.9×3.0 mm3.ATLAS 被划分为一个训练集（n=655）和验证集（n=300）.还选择了准确率（Precision）、F1 值、召回率（Recall）、交集度（IoU）这四个指标对相关方法的性能进行定量评估，等式（9）～等式（12）分别定义了这四个指标.

其中，TP 表示模型正确预测像素的数量，FP 表示该模型将阴性分类为阳性的像素数量，FN表示被错误分类为的正像素模型是负的像素数量.

3.2 实施细节

采用Pytorch 实现本文方法，模型的训练和推理均在一张NVIDIA GeForce GTX 2080Ti GPU 上开展，训练Batch 大小为12，优化函数是SGD，优化步长设置为0.9，权重衰减值lr=1e-4.此外，为提高模型的泛化性，防止模型过拟合，还对数据集进行了水平翻转、随机剪裁、旋转、对比度调整和添加随机噪声等数据增强操作.公式（6）中的参数α、γ、δ分别设为1.3、0.42、1.其他比较方法的相关参数按照原文模型的最佳性能被设定.

实验输入图像是MRI 的横向断层切片图像，每个断层切片图像的分辨率为224×224×3.在Swin Transformer 分支网络中，图像块大小被设置为4 个像素，Window 的大小为6，每次位移量为2.在该分支网络中，STrans Block 中的W-MHSA 和SW-MHSA的数量分别被设置为4，8，16，32.

在基于CNN 分支网络中每个Conv Block 中的卷积核大小统一设置为3×3，以及每个Conv Block中的卷积核数量分别设置为32、64、128、256、512.在每个卷积操作后，一个2×2池化操作也是应用于特征映射.在解码器块中，上采样层用于增加分辨率特征映射为2 倍，然后输出为与前一阶段的输出连接.最后一个输出通过使用1×1 卷积运算计算以及一个sigmoid 激活函数预测病灶区域概率和输出掩码图.

3.3 性能比较

本文选择了具有几个代表性分割网络（Unet［2］，Attention Unet［3］，Trans Unet［29］，Swin Unet［30］，MedSegDiff［34］，SegNext［35］）与本文模型进行比较.前面四种方法都采用了医学图像分割中常用的Encoder-Decoder 的U 形架构，其中Unet［2］是采用纯卷积方式提取输入图像的层次化表征，Attention Unet［5］在卷积操作的基础上加入了注意力机制.Trans Unet［29］是在Encoder-Decoder 分割模型基础上使用Transformer 构建输入图像的长距离依赖关系，Swin Unet［30］则是Trans Unet［29］基础上运用Swin Transformer从而让Transformer中的多头注意力模块能适应不同尺度的目标对象.MedSegDiff［34］是近年来流行的扩散概率模型在医学图像分割领域一次应用，该方法将分割目标看作是扩散概率模型要生成的对象，并在此过程中提出了特征频率分析器（FF-Parser）消除高频噪声对目标的影响.这些模型的参数是按照原始文章进行配置.SegNext［35］是在Encoder-Decoder架构中，采用了一种大核注意力机制（Large-Kernel Attention）获取目标的多尺度特征，并结合多层感知机（Multiple-Layer Perception Machine）对图像进行语义分割，相较于Transformer 模型有较好的推理效率.

为了验证不同方法在该数据集上的性能，请有经验的放射科医生将验证集按照病灶面积占图像背景的比值进行分级（见表1），表2显示了不同方法在不同层次验证集上分割结果的均值统计结果.总体上，采用 Encoder-Decoder 架构模型比MedSegDiff［34］效果更好，这说明该类架构对脑卒中病灶这类具有形状和尺度多变，前景与背景数据不平衡的情况依然具有优势.在Medium 组和Small组，本文方法相较于其他方法在各项指标中取得更好的效果.我们还观察到相较于纯卷积模型（Unet［2］，Attention Unet［3］，SegNext［35］），基于自注意力机制的模型（Trans Unet［29］、Swin Unet［30］）在相关指标值更差.这说明针对该数据集，较基于纯卷积的分割模型，基于Transformer 的分割模型在构建全局特征时更容易将背景信息引入病灶的真实区域，从而造成Transformer模型容易分割错误区域.

表1 验证集按照病灶面积占图像背景的比值分级Tab.1 The number of images at different levels for the four compared methods

表2 不同方法在不同层次验证集上分割结果的均值统计结果Tab.2 Quantitative comparison of the compared methods for the four metrics

图5 显示了不同方法的部分分割示例，第1 列是数据集中脑部T1-Weighted MRI 扫描数据中一份原始断层切片，第2列是手动标注病变区域的结果，其余各列则分别显示了不同方法在相应原始切片上的分割结果，这些分割结果按照病变区域的大小降序排列.从图5 显示可观察到前4 行所有方法都能较为正确的预测到病灶的正确位置.但与其他方法相比，本文方法的分割结果更接近在病灶边缘，从前4 行的最后1 列可看出，Swin Unet［30］的分割结果不够光滑，不能很好地保持病灶区域的形状.此外，当病灶区域面积较小时，本文方法依然能够预测病灶区域，尤其是最后两行待分割病灶的区域很小且形状难以辨认，这种情况对分割模型甚至人来说是非常困难的.即使如此，本文方法仍然正确地预测出了病变区域并较好的保持病灶区域形状，这进一步证明了本文网络分割难样本的能力.图6 显示了在相同处理过程中一些图像的特征映射CNN分支机构和Swin Transformer 分支机构.我们可以看到，基于卷积的局部特征和基于自注意力的全局特征是互补的.

图5 相关方法分割示例Fig.5 Segmentation example of the compared methods

图6 不同方法平均分割时间比较Fig.6 Comparison of average test time of different methods

4 讨论与分析

4.1 消融实验

本文对本模型中的Fusion Block 进行了消融实验，以分析该模块在本模型中的作用.具体方式是自顶向下逐渐减少Fusion Block 的数量，将STrans Block 和Conv Block 的输出特征图拼接之后直接输送至不同层次的Upsampler中进行解码.表3显示的结果可看出在同一验证集上，简化后的模型随着Fusion Block 数量的减少，其分割性能也逐渐下降.这表明在解码阶段对来自不同层次的局部特征和全局特征进行融合有助于提高解码器对病变区域的分割性能.

表3 融合模块对分割性能的影响Tab.3 Effect of the number of fusion blocks

4.2 模型深度的影响

本模型实验默认编码器有4 个处理单元.我们通过增加处理单元的数量，以验证网络的深度对分割性能的影响.表4 的结果显示当处理阶段的数量增加时，网络的分割指示会更好.但为了平衡模型的性能和训练效率，文中所有的实验都是采用4 个处理单元来验证本网络的性能.

表4 模型深度对分割性能的影响Tab.4 Effect of the network depth on segmentation performance

4.3 图像块大小的影响

图像块大小可改变Swin Transformer 分支网络的输入序列的长度.本文图像块Patch 的缺省值4×4，本文改变Patch 的大小对Transformer 分支网络性能的影响进行分析.表5 说明了本网络对不同大小的Patch 对最终分割性能的影响.由于图像块大小的平方与由输入序列的长度成正比，我们可观察到较小的图像块会带来较高的分割性能.可见较长的输入序列可使Transformer 能够在每个元素之间构建了更复杂的依赖关系，但同时也会增加模型训练的难度.

表5 图像块与输入序列的消融实验Tab.5 Ablation study on the patch size and the sequence length

4.4 效率分析

图6给出了不同方法在同一测试集中对每张图像分割的平均时间比较.从该图可看出，MedSegDiff［34］平均耗时最高，基于纯卷积的方法（Unet［4］、Attention Unet［5］、SegNeXt［35］）的效率优于本文所提方法.表6 显示了不同方法针对同一训练集训练后的参数（Params）大小以及浮点运算（GFLOPs）的比较.从该表可看出除了SegNeXt［35］和Swin Unet［30］，本文所提方法具有相对较少的可训练参数量和浮点计算量，总体上可以更好地满足临床对快速图像分析的要求.

表6 不同方法的参数量和浮点运算量比较Tab.6 Comparison of Parameters and GFLOPs on the different methods

4.5 Loss函数性能分析

我们使用本文的Loss函数（公式（8））分别对模型训练收敛性以及分割性能影响进行了分析.图7（a）显示了公式（8）作为Loss 函数在训练初期（约40 个Epoch），Loss 值下降较快，大约经过120 个Epoch，Loss 曲线变得较为平缓，模型逐渐收敛.此外，我们还使用F1值来衡量公式（8）与不同Loss函数对模型性能进行分析，从公式（11）可知F1值越大说明模型的稳定性和泛化性越好.如图7（b）所示，随着Epoch值增加，不同Loss 函数的F1 值也在上升，但本文的Loss 函数曲线所对应的F1 值始终是超过其他Loss函数，说明本文所使用的Loss 函数能够给模型带来更好的性能稳定性.