基于三重注意力机制的U-Net新冠肺炎肺部图像分割①

吴一锋, 于 瓅

(安徽理工大学计算机科学与工程学院,安徽 淮南 232001)

0 引 言

近年来,卷积神经网络(convolutionalneuralnetwork,CNN)在计算机视觉领域受到广泛运用,为图像的特征提取提供了更多的解决方案[1-2]。Long等人(2015)[3]提出了全卷积网络(fullyconvolutionnetworks,FCN),可以将任意尺寸的图像输入到模型当中,对最后一次获取的特征图利用反卷积操作进行上采样,使其图像尺寸与输入图像一致,便于对每个像素点进行预测,实现了端到端训练的神经网络,但是该模型只对最后一层特征层上采样,没有考虑低级特征与高级特征的关系,在训练过程中会损失部分特征。Ronneberger等人(2015)[4]提出了U-Net网络模型,该模型在医学图像领域得到广泛运用,在训练样本较少的情况下,该模型也能表现出较为优越的性能,完成图像分割的任务,但是对于病变的细微特征关注度不高。Zhou等人(2018)[5]提出了U-Net++网络模型,该模型对U-Net的跳跃连接进行了优化,使连接处的细粒度信息得到提升,但该研究需要大量CT图像。宋瑶等人(2021)[6]提出了基于U-Net改进模型的自动分割方法,在编码器中加入了EfficientNet-B0网络,增强对有效特征的提取,在解码器部分将上采样替换为DUpsamplin结构,以此提高病灶边缘细节特征的获取,但是对于小目标病灶识别度不高。谢娟英等人(2022)[7]提出了即插即用的融合多尺度特征的注意力模块,将其插入U-Net网络中,将复杂的跳层连接和通道间信息交互连接,捕获丰富的上下文信息,强化了网络的分割性能,但网络结构复杂,缺乏实时性。以上对于U-Net网络模型的改进,都能够在一定程度上提高目标任务的分割性能。

针对新冠感染的病灶分割问题,将经典的U-Net网络作为基本框架,对TA模块(TripletAttentionModule)进行改进,使其可以自适应选择卷积核大小,在跳跃连接处加入了改进后的模块,通过实现通道维度与空间维度之间的跨维度交互,将三个分支获取的精细张量进行汇总以进行自适应特征细化。

1 相关模型概述

1.1 U-Net网络模型

U-Net网络的提出一定程度上解决了医疗数据集相较于自然数据集较少的问题,该模型的端到端学习可以利用少量样本也能达到较好的输出效果。U-Net采用了编码器-解码器的结构,每个卷积层形成的特征图与对应的上采样层形成跳跃连接,从而增强对病灶区域的特征图获取,整体呈现U型结构。与FCN相比,U-Net没有直接在高级特征图中进行训练,而是在结合低级特征图后再进行上采样,使得最终获取的特征图同时具有低级特征和高级特征的信息,进一步提升了分割效果,但是网络经过池化层后无法还原病灶区域的细节信息,导致输出图像边缘处的病变与标签信息差距较大,对小区域病灶识别精度较低,难以满足病灶分割的需求。

1.2 视觉注意力TA模块

人类感知中的注意力涉及选择性地集中在给定的信息中,同时忽略其他部分信息,源于此类研究,计算机视觉中引入了注意力机制,该机制有助于完善感知的信息并保留上下文,对关键信息进行关注。Li等人(2020)[10]提出了TA模块,该模块由三个平行的分支构成,前两个分支负责实现通道维度C与空间维度H或W之间的跨维度交互,第三分支用于建立空间注意力,所有三个分支的输出都使用简单平均法进行汇总结构如图1所示。

图1 三重注意力模块

TA模块的作用原理是:将输入的特征图进行接收,将其传递给TA模块中三个分支的每个分支。在第一个分支建立了高度H和通道C的交互,先将输入的特征图沿H轴逆时针旋转90°,之后通过Z-pool层将其形状进行简化并输入标准卷积核,经过sigmoid激活层生成注意力权重,生成的权重被应用到旋转后的张量中,为保持与输入尺寸一致,将其沿H轴顺时针旋转90°。第二个分支与首个分支类似,由原先沿H轴旋转转换为沿W轴旋转,最终分支不对特征图进行处理,其他结构与第一分支一致。通过简单平均法将三个分支生成的细化形状张量进行汇总,以输出相同形状的精细张量。通过以上特征加强的方式,提高对特征的关注度。

2 本文模型

2.1 改进的TA模块

COVID-19病灶区域具有很多细节纹理特征,且含有小目标区域病灶。原有的TA模块Conv层为固定卷积核k大小的卷积操作,无法根据不同大小的通道维度C改变自身k的尺寸,由于一维卷积用于捕获局部跨通道交互,k决定了交互的覆盖范围,那么k可能需要随着不同通道维数而变化,尽管k可以手动调优,但这将消耗大量的计算资源[11]。其中的改进之一是将Conv层替换为自适应选择卷积核k大小的Adaptive-Conv层,通常来说,通道规模越大需要越大的卷积核,反之通道规模越小需要越小的卷积核,那么,k和C之间可能存在一种映射关系,即:

C=∈(k)

(1)

由于通道维数C通常设为2的整数次幂,那么根据给定的通道维数C,可以将卷积核k的大小自适应地确定为

(2)

其中|x|odd表示x的最近奇数,实验中将γ和b设置为2和1,即保证得到卷积核k的大小为奇数。

Z-Pool层是将旋转后的特征图分别进行全局最大池化(GMP)和全局平均池化(GAP),再将获得的张量进行拼接操作,全局最大池化用于保留更多纹理特征,全局平均池化则用于保留更多的背景信息,将两类全局池化划入两个分支,分别通过Adaptive-Conv层,最后再进行拼接操作。

设输入的特征图X=W×H×C其中W、H、C分别代表宽、高、通道维数,则g(x)表示全局平均池化,即:

(3)

那么全局最大池化可以表示为:

(4)

改进后的TA模块命名为AdaptiveTripletAttentionModule(自适应三重注意力模块),简称ATA模块,其结构图如下图所示:

图2 自适应三重注意力模块

如图所示,将输入的新冠感染CT图像特征图分为三个分支,保留了TA模块中旋转操作,使其进行跨维度交互,与TA模块不同的是,为了增强网络对局部细节特征的提取并提高对小目标病灶区域的识别度,便对全局平均池化和全局最大池化得到的张量分别输入Adaptive-Conv层,借用了残差结构将得到的特征图与输入后旋转的特征图进行拼接,能够有效提高网络的鲁棒性。

2.2 COVID-19病灶分割ATA-UNe模型

为了提高病灶分割的精度,以及针对U-Net网络的局限性,将其与ATA模块进行结合,使得编码器提取的特征图在经过注意力模块后与解码器进行拼接,即在U-Net网络模型的跳跃连接处添加ATA模块,加强网络在各个阶段对病灶区域细节信息的关注度,进而提升分割效果,ATA-UNet模型结构如图所示。

3 实验结果与分析

3.1 数据集

实验采用COVID-19的分割数据集[12],由专业的放射科医生使用MedSeg进行标注,数据集分为两个版次。第一个数据集包括100张CT图像切片,该切片来自大于40名COVID-19的患者,其中含有四种类型的真实分割蒙版,mask=1的是磨玻璃结节,mask=2的是肺部实变,mask=3的是胸腔积液以及mask=4的是背景,所有图像的尺寸为512×512。第二个数据集延伸至829张图像,来自9名疑似患者,其中373张已经被专业放射科医生评估为阳性,其余均为正常图像,该数据集中所有图像的尺寸为630×630,两个版次一共929张样本。

图3 ATA-UNet模型结构图

在对上述数据集中的掩膜图像进行整理并思考之后,决定只专注于分割COVID-19病灶区域,其中包含磨玻璃结节以及肺实变两种病变的CT图,并将掩膜图二值化,灰度值为胸腔积液的赋予0,其余除了背景均赋予255。为了增强数据对于训练网络需求的鲁棒性,采用随机裁剪、随机翻转、随机旋转等图像增强方法来扩充训练集当中的图像和掩膜。训练集包括1780张肺部CT图像和病灶分割掩膜,测试集包括140张肺部CT图像和病灶分割掩膜,将图像大小统一resize为512×512,并进行归一化处理。

3.2 评价指标

灵敏度(Sensitivity),为预测正确的病灶区域的像素数目占真实病灶区域中的百分比,即:

(5)

特异度(Specificity),为预测正确的背景的像素数目占所有真实背景的比例,即

(6)

相似系数(Dice),用于衡量两个样本的相似程度,取值范围为[0,1],Dice的值越高,说明两张图像相似度越高,分割效果越佳即:

(7)

平均交并比(mIou),计算分割区域重叠面积的比值,得出真值和预测值之间的相关度,即:

(8)

在医学影像领域,特异度和灵敏度是两个经常关注的指标,并将相似系数作为主要评价指标,其中TP为真阳性(即被预测正确的病灶区域的像素数目),TN为真阴性(即被预测正确的背景的像素数目),FP为假阳性(即被预测为病灶区域实际是背景的像素数目),FN为假阴性(即被预测为背景实际为病灶区域的像素数目),集合X和集合Y分别代表分割结果的像素和实际标签的像素,k代表类别个数,pii代表预测为i类别实际也为i类别的像素数目,pij代表预测为j类别实际为i类别的像素数目。

3.3 实验环境

实验使用Ubuntu18.04的操作系统,python版本3.8,PyTorch版本1.9.0,CUDA版本11,使用NVIDIA GeForce RTX3060显卡,显存14 GB。优化器选择Adam来更新网络模型权重,动量参数设置为0.9,初始学习率设置为0.001,batch-size设置为10,epoch设置为60,图片输入尺寸为512×512。

3.4 实验结果分析

实验中对比目前主流的U-Net模型,以及FCN和U-Net的改进模型,详细实验结果Loss和epoch关系如图4。

图4 各模型的Loss趋势图

从图4的Loss趋势图可以看出,各模型在epoch为50左右已经开始趋近收敛,ATA-UNet相较其他模型在epoch为15之后收敛速度更快,且Loss值基本趋于最低值。

为了验证ATA-UNet网络模型的性能,进行了多次对比实验,实验与FCN、U-Net、VGG-UNet进行对比,batch-size、epoch、学习率、激活函数等基本参数保证一致性,实验结果如表所示。

表1 不同模型对比

从表1可以看出,在U-Net模型中添加了ATA模块后,在病灶区域分割的Specificity方面,虽然提出的模型此类指标不是最优的,但是Sensitivity,Dice,mIou指标在一定程度上都有所提升。ATA-UNet比FCN,U-Net,VGG-UNet,TA -UNet在Dice上分别提高了4.02%,3.36%,3.19%, 0.68%,由此可知,ATA模块的添加有效加强了分割效果,另外,在Sensitivity和mIou上该模型也比U-Net提高了8.96%和1.54%,即对于背景的预测有了较大提升。综合所有指标数据,提出模型在对COVID-19病灶分割的任务中表现优异,能够更为准确分割出病灶部分。

图5提出模型与其他网络模型对比的可视化结果。从分割结果可以看出FCN分割效果一般,这是因为FCN没有考虑低级特征和高级特征之间的联系,只对最后一层特征图进行上采样,忽略了病灶区域之间的空间关系,导致细节信息丢失,对于小范围的病灶识别度不高,因此模型的分割结果精度较低。U-Net在小目标病灶的分割效果上比FCN更优异,但是对于病变的细微特征关注度不高,包含了大量的区域冗余,对于整体病灶区域较小的分割,其结果差距较大,容易受到与病灶相似的干扰区域的影响。VGG-UNet相较于U-Net提升不是太大,对细节特征提取有所提升,但不够明显。提出的模型对小区域病灶分割效果提升比较明显,基本与专家标注图一致,且对边缘信息的处理更为细致,即病变较为复杂的部分也能够达到准确的分割结果。综上,该方法在分割准确性和完整性上均有明显优势。

图5 不同模型分割结果

4 结 语

为了解决COVID-19病灶区域的分割问题,提出了一种基于视觉注意力增强的ATA-UNet模型,对TA模块进行改进,加入了根据通道数变化的自适应卷积核,在U-Net模型中加入了ATA模块,加强了对于病灶区域特征的关注程度,使得病灶图像复杂区域的边缘轮廓更加细致,小区域病灶分割准确性更高。在收集的COVID-19数据集上的实验结果表明,提出模型更具优势,为研究COVID-19病灶区域的分割提供了有效参考。在后续工作中,将优化模型细节,提高分割性能。