基于注意力机制的两阶段纵膈淋巴结自动分割算法
2021-03-07徐少伟秦品乐曾建朝赵致楷
徐少伟,秦品乐,曾建朝*,赵致楷,高 媛
(1.山西省医学影像人工智能工程技术研究中心(中北大学),太原 030051;2.中北大学大数据学院,太原 030051;3.山西医科大学第一医院,太原 030001)
(*通信作者电子邮箱zjc@nuc.edu.cn)
0 引言
纵膈是左右纵膈胸膜及其间所夹的器官和组织的总称,其中遍布纵膈淋巴结,而淋巴结转移是肺癌由原发灶向远处转移的重要转移途径。肺癌淋巴结的转移规律一般遵循肺内淋巴结、肺门淋巴结、纵膈淋巴结的顺序[1],因此纵膈区的淋巴结状况对于肺癌患者精确的临床分期、治疗选择和改善预后有着重要的参考和影响作用。国际肺癌协会(International Association for the Study of Lung Cancer,IASLC)根据纵膈解剖学结构将纵膈分为14 个区域[2],在肺癌检查中对每一个分区都应进行淋巴结采集,其中两侧肺癌主要对2、4至9等分区进行采样。因此准确检测纵膈区淋巴结并标注其对应分区对于医生诊断具有重要意义[3-4]。
计算机断层扫描(Computed Tomography,CT)影像具有密度分辨力高、横断层影像较清晰、敏感性高等特点,对纵膈淋巴结检查具有重要价值。根据医学先验,纵膈淋巴结主要存在于器官间隙,一般认为正常淋巴结直径在1 cm 以内,而直径大于1.5 cm 的淋巴结具有临床价值。在CT 断层中,纵膈淋巴结因形态为椭圆形,易与血管断面、左肺动脉层面主升动脉心包上隐窝、左右心耳、胸腺、胸内甲状腺等混淆;而且肿大纵膈淋巴结间尺寸差异大,小淋巴结识别困难。因此医生要辨别淋巴结及定位其所在分区非常困难,并且截至目前,CT图像纵膈淋巴结分割算法中评价指标Dice score 最高仅为52.3%[5]。目前医学CT分割算法主要是基于全卷积网络及其改进网络U-Net的,但是对于像CT序列这样的三维图像来说,使用三维卷积的U-Net(3D U Network,3D U-Net)分割精度会更高。深度学习分割纵膈淋巴结目前主要有两种方式:一种是集成不同二维分割网络分割[6]结果;另一种是加入纵膈器官信息的单阶段三维分割[5]。
医学图像淋巴结、结节等小目标以及淋巴结与周围软组织差异小对最终分割精度影响大,针对以上问题有以下几种解决方案:第一种思路是文献[6]提出的集成两种不同的分割网络如二维U 形网络(2D U Network,2D UNet)[7]和Mask-RCNN(Mask Region Convolutional Neural Network)[8]的结果。但由于仅集成了两种网络结果,并未针对不同尺度淋巴结进行处理,导致小淋巴分割精度低;而且该方法仅利用了序列图像的二维特征,所以最终分割精度和效率都不高。第二种思路是秦品乐等[9]提出的采用多阶段二维三维结合的思路,即对二维单张图像提取间隙后再分割来获得疑似淋巴结,并将疑似淋巴结层间合并为体积块后再输入三维卷积分类网络。该方法由于共采用了两次二维单张图像分割,血管、淋巴结单张图像差别小,假阳性高,易引入累计误差,之后再加入三维分类网络导致整体耗时长、效率低。第三种思路是文献[5]提出的通过非深度学习方法分割部分纵膈器官后,将分割淋巴结扩展为多分类任务,再向3D U-Net[10]中输入纵膈器官和淋巴结标注,能在一定程度上减少负样本,提高小淋巴结分割精度。第四种思路如文献[11]中通过第一步粗分割得到目标感兴趣区域块后输入细分割网络,粗分割和细分割网络都以3D U-Net为骨架,其中细分割网络加入多种二维编码器对特征进行集成。相较其他三种思路,第四种思路精度更高,基于三维卷积更适合于序列图像任务,因此本文也采用这种思路。通过分析和医学先验发现,感兴趣区域可以通过以纵膈器官为掩码获得,能有效减少背景干扰,提升小淋巴结识别效率,增加纵膈淋巴结与周围组织差异,加速网络训练过程。因此,本文将第四种思路与第二种思路结合:在第一阶段分割纵膈器官、纵膈淋巴结,在保留淋巴结的基础上以纵膈器官为掩膜去除纵膈内外干扰器官组织,能有效改善纵膈淋巴结分割效果。
在第一阶段的基础上,由于去掉器官后仍留有部分纵膈组织和肺部肿瘤,两者都和淋巴结特征相似,因此如何只关注纵膈淋巴结,忽略周围软组织、肺肿瘤的干扰是另一个关键问题。目前很多算法都引入了自注意力机制,通过计算特征图、通道间信息来获取对应权值以引导神经网络关注目标区域,获取整体位置依赖。Hu 等[12]发现不同特征图对关键信息的贡献不同,因此可以通过学习的方式来自动获取到每个特征通道的重要程度,然后依据重要程度去提升有用的特征并抑制对当前任务用处不大的特征,类似于通道注意力[13]。Jaderberg等[14]发现神经网络池化下采样操作直接将信息合并会导致关键信息无法识别,提出了空间转换模块结构,用于指导网络学习目标的空间特性,类似于空间注意力机制。Wang等[15]引入一种Non-Local块,输出层的每个位置与整个特征图的依赖都计算一遍,获取全局信息,增强相似特征,是一种自注意力机制。由此,注意力机制已经证明可以很好地提取关键特征。本文引入全局聚合模块和混合位置、通道注意力模块并扩展到三维后,计算任意两个位置、深度、通道之间的依赖。本文算法在纵隔淋巴结数据集中的Dice score 达到了0.701 1,优于文献[6]中加入纵膈器官标签的单阶段分割算法和文献[5]中加入纵膈器官标签的并行二维分割算法。
针对纵膈淋巴结分割中纵膈淋巴结尺寸差异大、正负样本不平衡、与周围软组织及肺肿瘤特征差异不明显的问题,通过医学先验和深度学习结合,根据淋巴结只存在于组织间隙的特点,本文设计了两阶段算法模型,并引入注意力机制增强淋巴结特征,减少软组织、肺肿瘤干扰。
本文主要工作如下:
1)设计了一种新颖的多阶段分割结构,通过医学先验去掉干扰器官等负样本,增强淋巴结信息,能有效提升训练精度同时减少训练时间,最后可以根据纵膈器官和纵膈淋巴结定位所在分区;
2)改造原始3D 分割网络,针对纵膈淋巴结可能与软组织、肺肿瘤混淆的问题,引入了全局聚合模块,获取全局、局部以及深度特征;
3)针对纵膈淋巴结尺寸差异大、软组织和肺肿瘤干扰分割等问题,改造原始编解码块短接,加入双注意力模块,将通道注意力替换为自适应感受野模块并扩展注意力机制到三维,增强纵膈淋巴结空间形态特征,大大增强了网络对目标的分割能力。
1 相关工作
目前基于深度学习的三维医学影像分割算法按阶段分为单阶段和多阶段两类,按维度分为二维与三维两类。单阶段例如U-Net、3D U-Net、U-Net++[16]、Non Local U-Net[17],这类方法一般修改现有网络结构,加入注意力机制[18-19]或残差块[20]提升网络结果;第二类是多阶段方法,主要有Mask-RCNN、文献[11]中的级联算法和文献[9]中的级联算法等。
多阶段分割网络通过对目标的候选区域进行粗分割、筛选,然后再进行精细分割。多阶段中粗分割可能会引入累计误差,通常采用修改损失函数、增大候选区域等方法以获得更高召回率,降低累计误差。
深度学习视觉任务中注意力模拟大脑只关注感兴趣区域并忽略其他背景干扰,与多阶段算法人为干预提取感兴趣区域是类似的想法。注意力机制将可用计算资源偏向更有用的特征信息[13]。
综上所述,本文方法结合了目前深度学习中纵膈淋巴结分割任务中文献[5-6]单阶段加入器官标签和多阶段分割任务中文献[9]和文献[11]的思想,设计了双阶段分割算法。针对纵膈淋巴结正负样本不平衡和尺寸差异问题,本文与文献[5-6]中方法相同,在阶段一中加入器官标签,同时分割淋巴结和器官;得到分割结果后经过数字图像处理,根据医学先验去掉原图的器官等干扰因素来提高阶段二网络分割精度并降低训练难度。针对纵膈淋巴结尺寸差异、软组织和肺肿瘤特征相似等问题,阶段二与文献[5-6,9,11]中方法的不同在于,针对纵膈淋巴结尺寸差异、软组织和肺肿瘤特征相似等问题,阶段二提出基于注意力机制的分割网络,引入并修改了全局聚合模块和双注意力模块来解决以上问题,进一步提升了最终分割精度。
2 基于注意力机制的两阶段分割算法
本文算法总体框架如图1 所示:第一个阶段采用3D U-Net 分割纵膈器官、淋巴结;第二个阶段将原图去掉纵膈器官后输入到注意力分割网络中进行精细分割。
图1 本文算法总体框架Fig.1 Overall framework of the proposed algorithm
2.1 获取纵膈淋巴结
根据医学先验,纵膈淋巴结存在于纵膈器官间隙间软组织中,纵膈处器官分别为食管、肺干、主动脉弓、升动脉弓、降动脉弓、奇静脉、心脏、静脉腔、头壁静脉、脊柱、肺静脉、锁骨下和颈动脉、肺、气道。其中心脏、肺部、脊柱、降动脉等器官尺寸较大或位置固定,特征明显,分割精度高。阶段一中,分割以上器官加纵膈淋巴结,其中器官分割平均Dice score 为0.869 0。得到分割结果后对整体器官做闭操作、填补空洞,得到整体胸腔内部,排除胸腔外软组织干扰,并且同时将脊柱排除,如图2(b)所示。之后在原图基础上保留淋巴结粗分割,以各器官为掩膜,去除器官后进行开操作以去掉毛细血管、过小结节等,如图2(d)所示。第一阶段处理后理想情况应包含全部淋巴结,但实际情况下经过第一个阶段处理后总体召回率为0.910 2,其中部分病例数据Dice score 和准确率可达0.50 以上,因此权衡引入的累计误差和阶段二中分割性能提升,认为可以接受第一阶段处理导致的累计误差。
在阶段一获取到纵膈组织间隙,原图包含纵膈淋巴结、血管、肺泡、其他结节和软组织。纵膈淋巴结单张图像与血管差别较小,采用三维卷积获取层间上下文信息可以在序列图中将两者区分开。纵膈淋巴结尺寸差异大,通过去除纵膈器官加滑动窗口的形式减少反例、背景,缩小输入尺寸。经过以上处理后输入引入注意力机制的分割网络中。
图2 纵膈间隙提取过程Fig.2 Extraction process of mediastinal space extraction
2.2 基于注意力机制的分割网络
本文提出的基于注意力的分割网络模型整体结构如图3所示,网络结构和3D U-Net 相同,采用编码器、解码器结构。主干网络采用标准3D U-Net 编、解码块,只在编码、解码块中下、上采样阶段加入全局聚合模块,在编、解码短接中加入双注意力模块。模型在第二到第四个编解码块中加入全局聚合(下、上采样)模块,特征块同时进入编码块和全局聚合模块中,最后通过元素求和的方式累加,解码块同理。模型在第三到第五个特征图concate 操作前通过双注意力模块。基于注意力的分割网络主要包含两部分结构:全局聚合模块(Global Aggregation Block,GAB)和双注意力模块(Dual Attention Block,DAB),其中双注意力模块包含位置注意力模块(Position Attention Module,PAM)、通道注意力模块(Channel Attention Module,CAM)。与普通卷积局部操作不同,GAB 主要通过计算全局位置间依赖,增强相似特征,增大感受野,减少在下、上采样中最大池化、反卷积丢失的信息。混合模块主要通过自注意力来捕获特征图的任意两个位置的空间依赖性(这一点与GAB 相似),然后通过不同卷积层得到不同感受野下目标特征后获取通道信息,经过瓶颈层后对相应特征块进行通道增强。
图3 网络模型总体结构Fig.3 Overall structure of network model
2.3 全局聚合模块
由于医学图像成像质量差,纵膈结构复杂,容易导致纵膈淋巴结与淋巴结周围软组织特征相似度高,影响网络判断,因此需要设计模型能够较好地获取区分淋巴结与软组织。卷积操作只能获得淋巴结局部特征,纵膈淋巴结的三维特征明显,因此希望能够获得全局包括深度特征、增加纵膈淋巴结和周围软组织间特征差别。Wang 等[17]提出的非局部注意力模块Non Local能够很好地解决这个问题,本文也引入Non Local后扩展到三维,以获取全局以及深度间依赖关系。如图4 所示全局聚合模块主要分为三个部分:query(Q)、key(K)、value(V),即Q(B×DQ×HQ×WQ×CK),K(B×D×H×W×CK),V(B×D×W×N×CV),其中:
式中:Unfold(⋅)为将D×H×W×C中深度、长宽拉为一维,变成(D×H×W) ×C;QueryTransformsCk(⋅)为可以改变输入特征图深度、长宽的任意操作,如降采样则一般采用卷积核为3× 3× 3、步幅为2 的3D 卷积;Conv_1Cx为采用1× 1× 1 卷积核将通道变为Cx的3D卷积。
图4 全局聚合模块Fig.4 Global aggregation block
在得到Q、K、V三个向量后,再进行如下计算:
其中:A为空间依赖权重,A与特征块V进行矩阵相乘得到O。O为增强后的特征,维度为(DQ×HQ×WQ)×CV,经 过dropout、Fold(⋅)还原后,使用卷积核为1× 1× 1改变O通道维度。最终编解码块特征图与O逐像素相加。
全局聚合模块可以减少最大池化造成的信息丢失,并且可以通过计算三维特征块每个位置间依赖关系,增强纵膈淋巴结三维特征,较原有网络能大幅提升纵膈淋巴结分割精度。
2.4 双注意力模块
纵膈淋巴结分散在纵膈器官间隙间的14 个分区中,为了将纵膈淋巴结分割和解析到与语义类别相关的不同图像区域,以及进一步区分软组织、淋巴结特征,引入自然图像处理中用于场景分割的位置、通道注意力模块。与场景分割问题相似,同样需要区分一些易混淆的类别,例如自然图像中“田地”和“草地”的区域通常难以划分。本文将双注意力分割网络(Dual Attention Network,DANet)[21]中两种注意力模块卷积、变换扩展为三维,将通道注意力改为自适应感受野后嵌入到编解码器同层连接中,在同层连接中,需要捕获任意两个位置间的空间依赖性,这样具有相似特征的两个位置都可以贡献出改进。与GAB 相比,同层连接中不需要对特征图尺寸进行变化,没有上下采样操作导致信息丢失,采用位置注意力模块与通道注意力模块更优。
其中,位置注意力模块实现原理类似全局聚合模块,局部特征A(B×D×H×W×C)经过三个卷积层,得到特征图A1(B×D×H×W×C)、A2(B×D×H×W×C) 和A3(B×D×H×W×C)。如图5(a)所示,分别经过Unfold(⋅)函数后,有
得到空间注意力权重S(D×H×W)2,它表示三维特征块中任意两个位置之间的依赖,两个位置越相似,相关性就越高。扩展到三维,可以加入深度的依赖信息,加强纵隔淋巴结蚕豆状特征。
最后计算得到最终结果E1:
其中超参数α初始为0,并在之后学习分配更多权重。位置注意力模块计算了所有位置间依赖权重,依赖权重与原始特征之间加权和可以得到增强特征。
图5 三维双注意力机制示意图Fig.5 Schematic diagram of 3D dual attention mechanism
文献[21]中通道注意力简单地将原始图像变为一维,忽略了纵膈淋巴结尺寸差异。本文为了能够得到不同尺度目标特征,修改文献[21]中通道注意力为自适应感受野[22],并扩展到三维,为了减少计算和显存,调整原结构中固定5× 5× 5卷积为级联3× 3× 3卷积和空洞卷积。由于文献[23]认为在浅层中级联卷积效果不佳,因此在浅层中替换为空洞卷积。由于空洞卷积容易引起网格效应,首先利用空洞率为2 的空洞卷积操作增大感受野,再级联空洞率为1 的标准卷积消除网格效应[13]。如图5(b)所示,自适应感受野模块计算分为分支、融合、选择三个阶段。
1)分支阶段:修改原有感受野的卷积,对于给定的特征图X采用3× 3× 3 卷积和级联3× 3× 3 卷积或空洞卷积两种不同感受野的卷积核构建两条不同感受野下的特征块和
2)融合、选择阶段:与文献[22]中方法相似,融合阶段首先通过逐元素求和操作获得融合特征U,通过全局自适应池化获得道信息S。引入文献[12]中的思想构建瓶颈层,对全局通道信息S进行降维操作,得到紧凑特征Z(d)。对紧凑特征Z进行升维操作、分离,得到a(c)、b(c)。紧凑特征Z的通道d由输入特征的通道数C、降维比例r及下界L决定,r、L默认为4,d=C/r且最小不低于L。
3)选择阶段:通过通道注意力权重a(c)、b(c)分别对两种不同感受野特征进行通道加权后逐元素求和,得到最终的融合特征E2。
通过加入自适应感受野模块,全自适应感受野网络比原分割网络拥有更丰富的不同尺度目标的语义特征,极大地提升了提取纵膈间隙、分割淋巴结的精度。
最后将通道注意力和位置注意力做逐像素加得到最终双注意力O=E1+E2,最后与解码器模块特征图进行concate。
3 损失函数
纵膈淋巴结分割任务为像素级进行分类任务,针对小目标前背景不平衡和医学分割等问题,本文损失函数定义为交叉熵损失、Dice 损失、Focal 损失[24]总和。总损失函数Ltotal定义为:
其中:LCE为交叉熵损失函数;LDice为Dice 损失;LFocal为Focal损失。
交叉熵损失定义为:
Dice损失函数定义为:
其中:p∈[0,1]n表示具有n个网络输出的预测值组成的一维向量,n为预测的像素数量;y∈{0,1}n表示二值真值标签向量,正类为1,负类为0。这里Dice不取正数,按原Dice计算。
Focal损失函数定义为:
4 实验与结果分析
4.1 实验数据集
实验数据集来源于公开纵膈数据集[7]与山西省肿瘤医院CT 放射科提供的纵膈区域CT 序列图像,图像分辨率为512×512像素,山西省肿瘤医院提供的数据集层厚为0.625 mm,由山西省肿瘤医院4 位经验丰富专家进行手工标注。放射科医师对照公开数据集标注的纵膈组织器官、淋巴结尺寸进行标注并多次校对。最终提取可用公开数据集和多位专家共同标注的淋巴结,获得60 例,共计2 800 张图像,其中包含1 310 张纵膈淋巴结图像。
4.2 模型训练方法和评价指标
实验数据随机抽取80%用于模型训练,剩余20%作为测试数据。数据预处理阶段采集训练集中标注的前景体素,基于值的0.5%~99.5%进行强度值裁切,收集训练数据间距,对训练数据采用三阶样条插值重采样。第二阶段滑动窗口大小设置为80 × 160 × 160,数据增强采用Batch Generators框架针对三维块进行随机镜像翻转、随机角度旋转、随机裁剪、弹性形变、随机缩放、伽马增强等方式增广。
本文模型采用Adam 作为随机梯度下降优化器,初始学习率为3× 10-4,学习率衰减为3× 10-5,每当训练损失的移动平均值在20个epoch内没有改善时,学习率降低0.2。训练使用一台NVIDIA P100 GPU 服务器,每个epoch 训练时间为637.78s,总计约70个epoch。
训练完成后本文选用查准率P(Precision),查全率R(Recall)和Dice score 共同评估本文模型和其他模型的性能。同时针对不同提取感兴趣区域方法、引入的注意力模块以及修改后的注意力模块结构进行了对比实验,验证本文方法合理性。其中Dice score计算公式如下:
其中:X表示为标签值;Y表示为预测的分割结果。
4.3 对比实验
首先,将本文提出的算法和其他常用的医学影像分割算法、纵膈淋巴结分割算法进行性能对比实验。数据集统一采用纵膈淋巴结数据集,基线方法采用标准3D U-Net方法,损失函数全部采用本文定义的损失,不同方法对比结果如表1 所示。相较其他算法,本文算法在各个评估指标上均有明显提升,并且各个性能指标均高于0.7,在纵膈淋巴结分割领域,本文算法达到了最高水平。不同检测算法的损失对比如图6所示。
表1 不同方法实验结果对比Tab.1 Comparison of experimental results of different methods
图6 不同检测算法的损失对比Fig.6 Comparison of loss in different detection algorithms
4.4 消融实验
为了验证注意力模块结构设计的合理性,本文设计了六种不同结构:第一种是直接使用3D U-Net,不加取间隙、无注意力模块;第二种取间隙后采用无注意力模块的3D U-Net;第三种只在编码、解码器上下采样块中加入全局聚合模块,去掉双注意力模块;第四种是仅在编解码短接中加入文献[21]中的位置、通道注意力模块;第五种是仅在编解码器短接中加入本文提出的双注意力模块;第六种为本文方法,加入全局聚合模块和本文提出的位置、通道注意力混合模块。根据表2 所示的结果可以看出,本文设计的基于注意力机制的分割网络可以更好地提升网络的性能。
表2 注意力模块实验结果对比Tab.2 Comparison of experimental results of attention modules
4.5 测试结果展示
不同方法分割结果的可视化如图7 所示,比较的方法分别为标签、基线方法、2D U-Net+Mask RCNN[5]方法、3D U-Net 加入组织标签(Structures)[6]方法和本文方法。其中第一行为不同方法对小目标和肺肿瘤等区分能力比较;第二行为不同方法对较大目标分割能力比较;第三行为不同方法对中等偏小较易分割淋巴结的分割能力比较。综上基本涵盖了纵膈淋巴结分割中的不同情况,可以看出本文方法分割效果更好,并且有效地解决了纵膈淋巴结与软组织、肺肿瘤难以区分的问题。
图7 分割结果对比Fig.7 Comparison of segmentation results
5 结语
本文针对目前深度学习分割纵膈淋巴结精度差、纵膈淋巴结分割难等问题提出了一个三维双阶段淋巴结分割算法,并设计了基于注意力机制的分割网络,主要用于纵膈淋巴结分割,并根据纵膈器官和淋巴结进行分区定位。首先,进行纵膈淋巴结粗分割与干扰器官去除操作,用于解决正负样本不均衡、纵膈淋巴结尺寸差异大问题;然后,引入了全局聚合模块、双注意力模块计算三维块整体信息和不同感受野下的特征信息,增强淋巴结特征,减少软组织、肺肿瘤干扰;最后,可以进一步结合纵膈器官、淋巴结来得到纵膈淋巴结分区。本文算法在纵膈淋巴结数据集分割上实现了0.726 9 的召回率和0.701 1 的Dice score,并且本文算法灵活简单,同样适用于其他淋巴结分割任务中。目前只使用了标准3D U-Net作为骨架,使用V-Net 等性能更优的分割网络作为主干将会更好地提升模型性能。因为受到显存限制,无法将全局聚合模块和双注意力模块用于第一、二层特征图,之后将针对两个模块作进一步改进,尝试加入搜索窗。