基于注意力增强U-Net的脑卒中病灶分割
2022-08-16王一诺张俊然李家琛
王一诺,张俊然,刘 彦,李家琛
(四川大学 电气工程学院,四川 成都 610065)
0 引 言
脑卒中是常见的脑部疾病,我国总体卒中终身发病风险已高达39.9%[1],其中急性缺血性脑卒中约占卒中总数的70%[2]。在缺血性卒中急救过程中,医生主导的阅片诊断环节关系着能否在4.5小时“黄金时间窗”内尽快恢复缺血脑区供血,降低致死致残率[3],但该过程易受医生经验等主观因素制约。利用计算机辅助诊断技术,可以准确分割影像中的缺血病灶,节约医生手动勾画的时间和精力,减少不同医生个体的主观影响。
区别于传统机器学习方法依靠人工设计特征,基于卷积神经网络(CNN)的分割模型可以更好地适应分割任务,自主学习特征,逐渐成为医学图像分割领域的主流方法。U-Net[4]是经典的2D分割网络,但是在对图像特征的关注上仍有不足。有学者提出了3D CNN方法实现缺血卒中病灶的分割,但网络级联[5]或增加网络深度[6]的方法会大量增加分割时间,不利于临床急救的开展,同时对GPU等硬件条件有更高要求。Wang等[7]利用CTP影像合成病灶特征更明显的伪DWI影像再进行分割,将3部分网络联合进行端到端训练,虽然适应了特定的分割任务,但也存在可移植性较差、单个网络调试困难等问题。已有研究表明注意力机制可以使CNN关注图像中对当前任务更关键的通道、空间等特征,从而提升网络性能,但大多是在物体识别的背景下提出的[8-10],或只涉及增强分割网络空间注意力的讨论[11]。
基于上述问题,本文提出了基于注意力增强U-Net的缺血性脑卒中病灶分割方法,主要创新点如下:
(1)为了适应病灶分割任务,提出了通道和空间注意力CSA(channel and spatial attention)模块,利用一维卷积和膨胀卷积,分别获得高效的通道间依赖关系和感受野更广的空间注意力,通过实验论证了CSA结构的合理性。
(2)为了增强U-Net对缺血性脑卒中影像中特征通道和空间信息的关注,将CSA模块与U-Net相结合,改善了U-Net的分割性能,并与其它U型网络的分割结果进行了对比,论证了本文模型的优越性。
1 基于注意力增强U-Net的缺血病灶分割
本文采用ISLES(ischemic stroke lesion segmentation)2018[12]数据集中的多模态脑影像作为实现缺血性脑卒中病灶分割的数据源,算法的主要流程包括数据处理、基于注意力增强U-Net的图像分割两部分,如图1所示。
图1 本文算法流程
1.1 实验数据介绍
ISLES 2018数据集提供了94例训练集和63例测试集影像数据,仅公开了训练集病灶标签。本文提出的算法在训练集的CT(computed tomography)和CTP(computed tomography perfusion)影像的4个灌注模态:脑血流量(cerebral blood flow,CBF)、脑血容量(cerebral blood volume,CBV)、平均通过时间(mean transit time,MTT)、达峰时间(time maximum,Tmax),及对应的病灶标签上进行验证,使用的各模态影像如图2所示。在每例数据中,水平面的2D图像大小均为256×256像素,根据病例不同像素间隔在1 mm左右不等,同时每个病例包含2至22张2D图像。
图2 患者的5个模态影像以及病灶标签(水平面)
1.2 数据处理
本文使用的影像数据特点可能对网络的训练过程产生不利影响:①原始CT图像中具有高强度的颅骨区域;②像素数值范围较广,可能导致分割网络训练不易收敛;③缺血病灶区域仅占大脑组织区域的7%,类不平衡情况较为严重。同时,从对称的大脑半球影像中可以学习更多区分病灶和健康组织特征[13,14]。综合以上因素,对原始数据进行以下处理。
(1)颅骨剥离:4个灌注模态的图像中有效数据区域不包含颅骨,利用此特点生成掩模对原始CT图像进行颅骨剥离。
(2)对称模态增强:将各模态图像分别沿人体解剖学的矢状轴翻转(水平翻转),采用刚性配准拟合翻转前后的图像。
(3)图像归一化:使用Z-score标准化方法对图像进行归一化处理,用以下公式表示
(1)
其中,IMOD和I′MOD分别表示归一化前后的图像,MOD表示各个影像模态,μ、λ分别为IMOD的像素均值和标准差,归一化后的图像满足均值为0方差为1的正态分布。
(4)图像块采样:分别以病灶组织和健康组织的像素为中心提取64×64大小的图像块。
(5)数据增广:采用表1所示的数据增广方式,使以病灶组织为中心提取的图像块达到与以健康组织像素为中心的图像块相同的数量N。其中弹性形变参数α和γ分别是控制变形强度的比例因子和弹性系数。实验中设定N=1000,进行数据增广后病灶区域的占比提升到了17%,类不平衡情况得到了改善。
表1 本文采用的数据增广方法和对应参数
1.3 基于注意力增强U-Net的分割网络
大部分CNN存在共有的问题[15]:构成CNN的卷积层在不同的空间位置使用共享的权值,可能导致CNN缺乏空间意识,进而在处理不同位置、形状、大小的目标时表现不佳;此外,大量特征通道中可能存在通道冗余的情况。U-Net这一经典的语义分割网络同样具有上述问题。
为了使U-Net适应不同位置、形状、大小的缺血病灶的分割任务,增强分割网络对通道和空间信息的关注,在U-Net的长连接添加通道和空间注意力机制,使其具有处理空间和通道信息的特定结构,更加关注特征信息丰富的通道,同时在空间维度对目标区域给予更多关注,使网络分割性能得到改善。本文提出的分割模型如图3所示。
1.3.1 通道和空间注意力机制CSA
为适应缺血病灶的分割任务,本文提出了通道和空间注意力模块CSA,结构如图4所示,使特征图在通道和空间维度依次获得注意力增强。其中应用到了使用快速一维卷积获得高效的通道关系和膨胀卷积获取更大的空间感受野的思想,同时对通道注意力增强模块CA和空间注意力增强模块SA都添加了残差连接的结构。
图4 本文提出CSA模块结构
假设输入特征图为F,通道注意力机制CA形成的加权输出特征图为F1,空间注意力机制SA形成的加权输出特征图为F2,则此注意力增强的过程公式为
F1=MC(F)⊗F+F
(2)
F2=MS(F1)⊗F1+F1
(3)
式中:MC(F)为特征图F经过通道注意力增强后的输出权值,MS(F1)是特征图F1经过空间注意力增强后的输出权值,⊗表示特征图加权乘法。
(1)通道注意力模块
通道注意力模块CA可以增强分割网络对信息丰富的特征通道的关注,抑制冗余通道的表达,在ECA(efficient channel attention)的基础上添加了残差连接,由全局平均池化层、一维卷积层和权重计算层组成。其原理是首先对特征图的信息通过平均池化的方式在空间维度上进行压缩,然后通过一维快速卷积,利用相邻通道的信息对每个通道的注意力权重进行预测。其中不涉及SE(squeeze-and-excitation)和CBAM(convolutional block attention module)中通道注意力机制广泛使用的通道降维方法,避免了降维给通道注意力预测带来的负面影响和获得低效的通道间依赖关系。残差连接可以使特征映射对输出的变化更加敏感,防止网络发生退化[16]。
将平均池化和卷积过程、sigmoid归一化分别用Favg.pool、Fconv(X;inch,ouch,k,d)、 σ的形式进行表示,X、inch、ouch、k、d分别表示输入特征图、输入通道数量、输出通道数量、卷积核大小、膨胀率,其中膨胀率d在普通卷积层中默认为1。则通道注意力机制输出权重MC(F)的计算公式为
MC(F)=σ(Fconv(Favg.pool(F);1,1,k,1))
(4)
一维卷积层的卷积核大小k和通道数C满足指数映射,因此可以自适应地计算卷积核大小k,即
(5)
其中, |t|odd为取计算结果最接近的奇数。
(2)空间注意力模块
空间注意力模块SA可以增强分割网络对前景像素的关注,抑制分割不相关区域特征的表达,由膨胀卷积层、全连接层和权重计算层组成,同样添加了残差连接。其原理是首先以卷积核大小为3膨胀率为d的膨胀卷积层学习特征,同时对特征通道进行降维,并使用ReLU函数激活,最后通过全连接层对通道维度进行恢复,获得分布式特征在特征空间中的映射,最终通过权重计算层得到特征图进行空间注意力增强的权重。这样的设计与SE模块的挤压激励原理对比,以膨胀卷积层代替原有的全局平均池化层对特征图的压缩操作,在不改变卷积层的参数总量和输出特征图大小的基础上,增大了卷积核对空间层面特征信息的感受野[17],获得了更多可利用的上下文信息,大小为3的卷积核在不同膨胀率下的感受野原理如图5所示,从左到右膨胀率d依次为1,2,3。
图5 不同膨胀率下的感受野
空间注意力机制输出权重MS(F1) 的计算公式为
(6)
其中,C为特征图的通道数,r为通道压缩率,d为膨胀率。考虑到性能和计算成本,在实验中参数选择为r=4和d=3。
1.3.2 注意力增强的U-Net网络
在医学图像分割的深度学习方法研究中,U-Net的出现一定程度上改善了利用有限样本训练模型完成分割的问题。U-Net具有对称的编码器和解码器结构,在编码器部分,对特征图进行4次下采样得到高级特征图,再通过解码器部分进行对应的4次上采样,最终将得到的特征图恢复到原图片的分辨率,获得更精细的分割细节信息。此外,U-Net的另一个显著特征,在编码器每层次最终获得的特征图,会通过长连接传递到解码器部分,这样的结构可以将每层次获得的不同尺度的低级特征和高级特征融合,从而将低级特征作为高级特征的补充信息,用以修正并获得最终的分割结果,同时也避免了直接对高级特征图进行监督和损失函数的计算,提高了分割精度。
在本文的研究中发现其分割结果与病灶标签相比,得到的病灶位置、大小和形状都有一定的差距,产生了较多的误分割和漏分割,仍存在一定的改进空间。将U-Net与本文提出的CSA模块结合,在不会过多增加原有分割网络体积的基础上,使其对图像中的病灶区域和信息丰富的特征通道给予更多关注,抑制冗余通道和分割不相关区域特征的表达。如图3所示,将CSA模块作为处理通道和空间关注信息的结构,添加在U-Net的长连接, 以加强模型对不同尺度的低水平特征图中病灶的通道和空间关注度[18],从而加强对病灶区域特征的学习能力,能够对其进行更准确的定位,适应不同大小和形状病灶的分割任务,提高病灶分割精度。
2 实验结果与分析
本文实验在Ubuntu 16.04的Linux设备上运行,深度学习框架基于PyTorch,硬件配置为Intel®CoreTMi7-9700 @3.00 GHz CPU,内存为16 GB,显卡为NVIDIA GeForce GTX 2060。
2.1 评价指标
为了量化评估本文实验涉及到的网络模型的分割结果准确性和离散性,选择Dice系数、灵敏度(Sensitivity)、阳性预测率(PPV)、Hausdorff距离(HD)作为评价指标
(7)
(8)
(9)
(10)
公式中做出以下定义:正确分割的病灶区域为真阳性(true positive,TP),错分为病灶的区域为假阳性(false positive,FP),正确分割的非病灶区域为真阴性(true nega-tive,TN),错分为非病灶的区域为假阴性(false negative,FN),A和B分别表示病灶标签与分割结果中病灶边缘的点。Dice系数、灵敏度、阳性预测率都是可以从像素层面衡量分割结果中正确分割、误分割以及欠分割程度的综合性指标,但是对于分割结果中的离群值无法表征,因此引入了Hausdorff距离,用于评价病灶标签和分割结果边缘的差异性。
此外,为了对分割结果中病灶位置的准确性进行评价,本文引入了质心距离(centroid distance,CD)的概念。即以病灶标签f的质心(x1,y1)和分割结果g的质心(x2,y2)之间的距离评估分割模型对病灶位置预测的能力,用以下公式表示
(11)
2.2 实验结果与分析
本文使用1.1节介绍的数据集进行实验,将306组包含病灶的多模态图像按照4:1的比例分为训练集和测试集。采用自适应优化器Adadelta,训练批大小为32,损失函数使用交叉熵与广义Dice损失函数之和,采用dropout机制和早停法监督训练防止过拟合。实验结果表明,本文提出的CSA模块可以提升U-Net对缺血性脑卒中病灶的分割性能,能够更好地定位病灶。
2.2.1 消融实验
为了论证CSA模块结构对U-Net分割性能的提升,以U-Net为基础网络,进行了5项对比实验:①不添加注意力机制;②添加无残差连接的NCSA模块;③添加交换CA和SA级联顺序的SCA模块;④添加文献[19]提出的并联结构改进的PCSA模块;⑤添加CSA模块。对CSA模块的结构调整如图6所示。
图6 对于CSA模块的结构调整
分割结果评估见表2,在灵敏度和阳性预测率较为稳定的情况下,添加CSA后U-Net的Dice系数、Hausdorff距离、质心距离最优,在病灶分割精度和分割结果的离散性、病灶定位的准确性上表明CSA优于其它3种结构,验证了本文提出的CSA模块在结构上的合理性。
表2 调整CSA结构的分割结果评估(均值±标准差)
2.2.2 分割性能对比
为了进一步评估本文模型性能,选择了ISLES 2018挑战中前两名所提出的分割网络部分SL-Net[7]和不对称U型残差全卷积神经网络[20],以及同样在长连接处添加注意力机制CBAM的U-Net网络结构,将上述网络在相同实验条件下的分割性能进行比较,评估结果见表3。同时将病灶标签、U-Net、上述网络模型,共6组结果对比绘制,使用不同颜色表示过分割、欠分割、正确分割区域,如图7所示。
根据表3的实验结果可以得出:在与不同的U型网络的分割结果对比中,本文模型拥有最佳的Dice系数和质心距离,在分割精度和对病灶的准确定位上具有一定的优势。结合图7进行分析:在第一行分割结果中,本文模型较好地分割出了缺血病灶,不存在离散的假阳性区域;所有模型在第二行分割结果中普遍存在较大范围的过分割情况,本文分割模型对此进行了有效改善;第三行分割结果中,除本文模型之外,其它模型的欠分割情况较为严重,在临床应用中,不能被检测出来的病灶对病人的影响是不可预估的。虽然本文模型的各项评估参数并不全是最优,但是在正确分割、过分割和欠分割的综合性能,以及对病灶的准确定位方面,本文模型表现最佳。关于进一步减少假阳性分布的问题,后续研究中可以采取更合适的阈值处理预测结果,或根据CTP影像各模态反映大脑血液动力学变化的特点添加病灶区域的约束进行改善。
表3 不同U型网络的分割结果评估(均值±标准差)
图7 不同U型网络的分割结果对比
在不对称U型残差FCN的分割结果指标中Hausdorff距离最优,这一优势原作者在文中也提及,但从图7的结果来看,该模型的欠分割情况较为严重;SL-Net是文献[7]中联合训练网络的分割网络,在U-Net的编码器各层应用了通道注意力机制SE,相比于U-Net分割性能获得了一定改善,但整体结果不及本文模型;CBAM-U-Net作为同样使特征图在通道和空间维度依次获得注意力增强的对比实验,在正确分割、分割结果的离散性和病灶定位的准确性的综合评价上不及本文模型,表明本文提出的通道和空间注意力机制CSA更适合于缺血病灶的分割任务。
以单张图像的分割时间作为衡量本文算法时间效率的指标,本文模型与上述U型网络参数量与单张图像平均分割时间对比见表4。可以看出,本文模型对256×256大小图像的分割平均用时为0.73 s,相比于U-Net仅增加不到0.1 s,模型参数量增加不到3%,在可接受范围内获得了分割精度的有效改善;同时与其它3种分割网络对比,本文模型的单张图像分割时间相对最优。
表4 不同U型网络的参数量和单张图像平均分割时间/s
3 结束语
本文结合了一维快速卷积获得高效通道注意力和膨胀卷积扩大感受野的原理,提出了通道和空间注意力模块CSA,并将其与医学图像分割经典网络U-Net相结合,利用注意力机制使U-Net对各级提取到的低水平特征在通道和空间层面加以关注和利用,再与高水平特征进行融合,使分割网络对病灶区域与非病灶区域、区域边界的特征、特征通道信息进行关注进而充分学习。实验结果表明,本文模型在没有过多增加U-Net网络结构复杂性的基础上,提高了对多模态缺血性脑卒中病灶的分割精度,并且对于准确分割、误分割和过分割的综合性能较好,对病灶的定位较为准确,可以为临床医生对于缺血性脑卒中患者的诊断、治疗方案的制定、预后等提供客观高效的决策支持,同时本文实验方法为研究缺血性脑卒中病灶区域的快速量化评估提供了一种思路。关于通道和空间注意力机制CSA在目标检测、目标跟踪等其它CNN机器视觉领域的作用,未来可以进一步探究。