一种具有边缘增强特点的医学图像分割网络
2022-05-31孙军梅葛青青李秀梅赵宝奇
孙军梅 葛青青 李秀梅 赵宝奇
(杭州师范大学信息科学与技术学院 杭州 311121)
1 引言
近年来,深度学习技术已广泛应用于医学图像领域[1–3],其中如何自动识别和分割医学图像中的病灶是最受关注的问题之一。由于人体器官多样、病灶形状复杂、图像噪声干扰等诸多原因,器官病灶等一些待分割物体容易出现边缘分割不清晰、缺失值大等状况。
目前已有较多学者针对医学图像分割方法展开相关研究,其中UNet[4]是最典型也是应用最广泛的方法,它利用收缩路径获取特征信息,利用扩张路径实现精确定位,在各类数据集上均有较好的表现。其后出现了许多在UNet上进行改进的网络模型。Shankaranarayana等人[5]将残差连接思想和UNet结合提出残差UNet(Res-UNet)。Oktay等人[6]提出注意力UNet(Attention U-Net),通过集成注意力门捕捉显著特征。Zhou等人[7]提出UNet++,整合不同层次的特征,使用灵活的网络结构配合深监督,使得参数量巨大的深度网络能够在可接受的精度范围内大幅度缩减参数量。Jafari等人[8]提出密集残差UNet(Dense Residual U-net, DRU-Net),在残差网络(ResNet)[9]与密集卷积网络(DenseNet)[10]的基础上额外增加了跳跃连接,以更少的参数实现了更高的精度。此外,针对目前方法不能很好地分割肿瘤和器官的重叠部分这一问题,Li等人[11]提出了一种基于位置导向的可变形UNet,利用可变形卷积的空间变形能力处理器官和肿瘤的几何变换,并引入了一个新的池化模块,来保留传统最大池化操作中丢失的位置信息。为了自动完善分割网络的分割效果,Kitrungrotsakul等人[12]提出了一种用于医学图像分割的交互式深度优化网络RefineNet,由分割主干UNet和优化网络两部分组成,结合了多尺度初始细分的功能和用户的种子点。网络可以在训练时自行生成种子点,但在测试时需要用户提供种子点,较为不便,且随机误差较大。
医学图像通常对比度较低,组织特征具有可变性,不同组织之间或者组织和病灶之间边缘较模糊,且微细结构(如血管、神经等)分布复杂。图像边缘是医学图像最基本的特征之一,往往携带着一幅图像的重要信息。因此,边缘增强处理是医学图像处理的关键技术之一,对于进行高层次的特征描述、识别和理解等有着重大的影响。针对医学图像边缘增强问题,Zhang等人[13]提出了边缘注意网络(Edge aTtention Network, ET-Net),嵌入边缘注意表示来指导分割网络。具体地,边缘引导模块用于学习早期编码层中的边缘注意表示,然后将其传输到多尺度解码层,并使用加权聚合模块进行融合,对边缘信息做到了很好的获取。Valanarasu等人[14]提出了一个完整的级联网络KiU-Net进行脑解剖分割,可以同时利用Ki-Net的底层精细边缘特征图和U-Net的高级形状特征图,不仅提高了分割精度,而且对于小的解剖标志和模糊的噪声边缘也实现了快速收敛。针对分割区域及边缘的不确定性与模糊性,Lee等人[15]提出新颖的边缘关键点选择算法和结构边缘保持分割框架,使得模型能够自动提取和学习结构关键点信息,但在网络中增加了较多参数,不易于模型部署。针对区域边缘和区域内的不连续性,Chu等人[16]提出一种利用简单边缘检测器定位所有不连续点,并对这些区域进行额外监控的方法,有效提高了分割精度。
上述工作都为后续医学图像分割技术的研究提供借鉴意义,但仍有许多问题亟待解决。首先,分割精度较低。细胞和器官等形状复杂多样,且可能存在团簇重叠状态,致使错误分割的概率很大,影响医学判断。其次,模型结果的提升伴随着测试时参数的增加。越来越多的网络模型基于UNet做出改进[17–20],或采用双流卷积神经网络结构增强边缘信息[21–23],虽然在一定程度上可有效提升分割结果,但存在测试时参数冗余现象。
为了解决上述问题,本文基于UNet提出一种具有边缘增强特点的医学图像分割网络(Add-and-Subtract UNet, AS-UNet),Add是指训练时在UNet的基础之上增加边缘注意模块(Boundary Attention Block, BAB), Subtract是指测试时舍弃BAB。AS-UNet通过掩膜边缘提取算法得到掩膜边缘图,在UNet扩张路径的最后3层引入结合多尺度特征图的BAB,强化边缘,减少缺失值;并提出使用组合损失函数,结合基于区域的Dice损失(Dice Loss)[24]和基于边缘的Boundary损失 (Boundary Loss)[25],在保证区域缺失值小的同时补充边缘信息,提高分割精度;此外,由于组合损失函数包含了经过BAB前后的两种输出,在训练时经过前后向反馈不断更新AS-UNet中的网络参数,使得训练好的模型在测试时可以舍弃添加的BAB部分,保证分割精度的同时减少测试时参数。
2 具有边缘增强特点的医学图像分割网络AS-UNet
2.1 网络结构
本文所提具有边缘增强特点的医学图像分割网络AS-UNet结构如图1所示。采用UNet网络的编码—解码结构,输入通过UNet得到的直接输出称为输出1(output1)。将UNet扩张路径中L7, L8, L9得到的特征图R7,R8,R9分别作为输入层送入BAB。R7经过BAB后采用2× 2的上采样(Up-conv 2× 2)达到与R8相同的尺寸,作为补充层传入L8的BAB中。同理,R8经过BAB后采用2× 2的上采样达到与R9相同的尺寸,并作为补充层传入L9的BAB中。而最后一层得到的特征图R9经过BAB后则直接进行1× 1的卷积得到输出2(output2)。此过程如式(1)所示
图1 AS-UNet网络结构
训练时输入图像通过UNet生成输出1,同时将R7, R8, R9通过BAB还原图像边缘细节信息,生成输出2。将输出1和输出2分别与掩膜(Mask,即标签(Ground Truth, GT))进行损失计算,在前后向传播中更新参数,得到最终模型。
2.2 边缘注意模块
在UNet中,边缘分割模糊以及部分区域缺失仍是最主要的不足。因此,本文所提AS-UNet网络通过在UNet网络中引入边缘注意模块BAB以强化边缘信息,补充缺失区域。BAB结构如图2所示。
图2 BAB结构
在BAB中,改变掩膜边缘图大小并按通道连接至各特征图之后,有效地提供了精确的图形边缘,增强图像对于目标区域的边缘信息获取,使得后续网络更重视这一部分的信息,从而弥补UNet对边缘信息的忽视。
2.2.1 掩膜边缘提取算法
为了引入像素级别的精确边缘信息,本文通过掩膜边缘提取算法将训练图像的掩膜导出为掩膜边缘图,以此作为边缘信息的重要补充。掩膜边缘提取算法过程可表述为:遍历掩膜的每一个像素点(i,j),当遍历的像素点值为0,且以该像素点为中心的九宫格内其余像素点不都为0时,将该像素点记作0,直到掩膜的所有像素点都遍历结束后,生成掩膜边缘图。图3为掩膜与经过掩膜边缘提取算法导出的掩膜边缘图。
图3 图像掩模与对应边缘图
2.2.2 注意力模块
受空间通道压缩与激励(spatial-channel Sequeeze & Excitation, scSE)[26]注意力模块的启发,本文提出一种新的注意力模块,具体结构如图4所示。对于输入特征图U ∈ℜw×h×c,分别在通道和空间上进行压缩得到特征图UˆsCE∈ℜw×h×1和向量
图4 注意力模块
cSE和sCE顺着空间和通道维度分别进行特征压缩,将得到的特征权重向量分别加权到输入特征图上,完成不同维度上对原始特征的重标定。而scSE的工作仅仅是将cSE和sCE的输出结果逐像素相加,同时完成空间以及通道上的压缩与激励。而本文所提出的注意力模块首先将空间和通道上压缩得到的特征图进行相乘,得到与输入相同大小的权重W,如此便可以为输入特征图的每个像素都提供各自的权重,且相乘能够让重要的位置更加突出,而价值较小的位置则可以被抑制。随后再将这一权重与输入逐像素相乘得到最终输出。此方法相较于scSE同样简单且不增加任何参数,但能够更多地关注重要位置的信息,获取感兴趣区域及边缘特征信息。本文将其应用于边缘注意模块3×3的卷积层之后,帮助模型取得更好的分割效果。
2.3 组合损失函数
在图像分割领域,主要采用的损失函数包括基于分布的交叉熵损失函数(Cross-Entropy Loss)以及基于区域的Dice损失函数。交叉熵损失函数单独评估每个像素矢量的类预测,然后对所有像素求平均值,使得图像中的像素能够被平等学习。但是,医学图像中常出现类别不均衡的问题,由此导致训练会被像素较多的类主导,对于像素较少的物体很难学习到其特征,从而降低网络的有效性。而Dice损失函数,本质上是衡量两个样本的重叠部分,虽然解决了类别不均衡的情况,但是未关注到图像的边缘信息。由于医学图像的边缘信息尤为重要,因此本文提出结合基于区域的Dice 损失和基于边缘的Boundary 损失的组合损失函数L,在两个不同的侧重维度上进行监督。Dice损失和Boundary损失的定义如式(4)、式(5)所示
损失函数L结合了基于区域的Dice损失和基于边缘的Boundary损失,使得网络同时关注区域与边缘信息,在保证区域缺失值小的同时补充边缘信息,从而提高分割精度。随着神经网络的不断迭代,平衡系数α和β自学习调整更新,促使在前半部分的UNet网络中,Dice损失占据比重较大,相对来说更加关注区域的信息,在后半部分的边缘注意模块中,Boundary损失占据比重较大,相对来说更加关注边缘信息。本文使用组合损失函数发挥边缘注意模块的作用,达到关注区域信息的同时不损失边缘信息,解决目前医学图像分割中存在的分割缺失值大、边缘不清晰的问题。
BAB的引入提高了模型的分割精度,但也不可避免地增加了网络参数,本文方法AS-UNet在测试时利用神经网络结构的冗余性进行模型精简,只需要选择模型的输出1,达到不增加网络参数的目的。
这种方法可执行的依据是使用新的组合损失函数L,L中的Dice损失和Boundary损失都结合了y1和y2两种输出,因此,在将L总体最小化时必然会将针对y1部分的Loss最小化,如式(9)、式(10)所示。训练时经过前后向反馈不断更新AS-UNet中的网络参数,修复未留意的边缘信息。因此可实现在最终测试时即使舍弃边缘注意网络,仅采用y1作为测试结果,也不影响分割效果,在提高分割精度的同时极大地减少了模型参数,为实际应用提供了更大的可能性。
结合两种输出的损失函数L可以兼顾输出1与输出2,使其分别达到最好的效果。训练时输入需要经过基础网络+BAB,不断更新优化网络内部参数,而测试时可以舍弃BAB,在训练好的模型基础上选择输出1即可得到最终结果,在提高分割精度的同时可以减少参数量。
3 实验及结果分析
3.1 数据集
数据集1:Glas。Glas是MICCAI2015腺体分割挑战赛的公开数据集,由165幅图像组成,这些图像来自16个苏木精和伊红(Hematoxylin-Eosin,H&E)染色载玻片的大肠癌组织切片。原始图像大小不一,大多数为775×522。将数据集分为训练集和测试集,100幅图像属于训练集,65幅图像属于测试集。
数据集2:DRIVE。DRIVE发布于2003年,是一个用于血管分割的数字视网膜图像数据集,由40幅图像组成,其中7张显示出轻度早期糖尿病视网膜病变迹象。将数据集分为训练集和测试集,30幅图像属于训练集,10幅图像属于测试集。
数据集3:ISIC2018。ISIC2018[27,28]是用于黑色素瘤检测的皮肤病变分析数据集,其中用于医学图像分割的部分包括2594幅RGB皮肤病变图像,采用二进制标注。将数据集分为训练集和测试集,2074幅图像属于训练集,520幅图像属于测试集。
3.2 评价指标及实验细节
本文使用平均Dice(Mean Dice)[29]、F1值、Hausdorff距离[30]作为评价指标,衡量分割结果的优劣,更为公正地对分割结果做出评价。其中Mean Dice更多关注区域信息,用来衡量两组样本之间的相似性;F1值是精确率(Precision, P)和召回率(Recall, R)的加权调和平均,用来衡量单个样本目标的检测精度;Hausdorff距离更多关注边缘信息,用来衡量基于边缘的分割精度。3类评价指标计算方式如式(12)—式(14)所示。
实验使用的深度学习框架是Tensorflow 1.14,计算机操作系统为64位Ubuntu 20.04.1,同时使用GPU加速网络模型的训练和测试。采用Adam优化器,因为 Adam 算法可以在训练时自适应地调节学习率,且有更快的收敛速度,初始学习率设置为0.001。考虑到计算资源的局限性,训练阶段批大小(batch size)设置为2,测试阶段批大小设置为1,输入图像统一转换为512×512,并进行增强对比度处理,灰度化彩色图像后将灰度像素转化为0~1,减小输入特征的尺度,再送入网络模型进行训练和测试。
3.3 实验与结果分析
(1)不同模型对比实验。为验证本文方法的有效性,在3个不同类型的数据集上,分别和UNet及其他3个也在UNet基础上做改进的较新颖的图像分割方法UNet++, DRU-Net以及KiU-Net进行对比实验。其中UNet是医学图像分割的基础网络;UNet++在不同层之间重新设计跳跃连接将编码器和解码器相连,并且通过减轻未知的网络深度和设计一个修剪方案来提高模型性能;DRU-Net提出了一个简单但高效的网络模块,以更少的参数实现更高的精度;KiU-Net则基于UNet进行模块结构的改进,通过改变感受野捕捉精细边缘特征图和高级形状特征图,更好地实现边缘细节捕获与分割。
各模型的实验结果如表1所示。相比于其他4种网络,本文方法在不同数据集下3个评价指标均有较好的提升。与UNet相比,本文方法几乎不增加测试时参数量,在Glas数据集上Mean Dice和F1值分别提高了2.19%和2.21%,Hausdorff距离减小了31.8;在DRIVE数据集上Mean Dice和F1值分别提高了2.16%和2.64%,Hausdorff距离减小了10.89;在ISIC2018数据集上Mean Dice和F1值分别提高了1.53%和2.91%,Hausdorff距离减小了7.53。由此可见本文方法的有效性和较好的泛化能力。
表1 不同模型在不同数据集上分割结果对比
图5 迁移模型
图6—图8给出了5种网络模型在3个数据集上的分割效果。从中可以看出本文所提AS-UNet在分割结果中更加贴近原图的GT,尤其是在边缘的处理上更加细致。在Glas数据集中,细胞较大,细胞之间的粘合问题较少,本文方法处理得到的分割图存在的缺失值较少,且边缘较平整;在DRIVE数据集中,人眼血管较细,分布复杂,分割难度较大,本文方法处理得到的分割图相较于其他模型在细血管的连续性问题上处理得更好,其他模型分割得到的细血管很多都变成了点状而非连续,本文模型借助边缘注意模块在一定程度上仍能保持血管的基本连续;在ISIC2018数据集中,病变区域形状各异,各样本间差异较大,其他模型均存在明显的缺失或多余块,而本文方法较为精细地分割边界,且区域完整。
图6 不同模型在Glas数据集上的分割结果对比
图7 不同模型在DRIVE数据集上的分割结果对比
图8 不同模型在ISIC2018数据集上的分割结果对比
(2)消融实验。采用消融实验分别验证本文所提边缘注意模块BAB以及损失函数的有效性、可移植性和参数情况。所有实验训练配置均相同。参照图5,基础网络可用UNet和FCN等网络代替,因此分别以UNet和FCN为基础网络设计实验。其中,测试时,输出1可作为舍弃BAB后的直接输出,也可用于生成掩膜边缘图,继而送入BAB生成另一个输出2。将训练时添加BAB且测试时采用输出2的方法记作+BAB;将训练时添加BAB且测试时舍弃BAB采用输出1的方法记作+Sub。
如表2所示,各种对比方法在两种基础网络之上的效果相似。添加BAB后,通过引入掩膜边缘图解决分割边缘不清晰问题,并对网络关键特征的提取有一定的引导意义。相较于单纯的UNet和FCN,在各数据集上Mean Dice提升1%~4%,F1值提升1%~3%,Hausdorff距离也基本上有所减小,可见本文设计的BAB具有可移植性与非依赖性;对比+BAB与+Sub,两者在训练时都融合了BAB,但在测试时,前者使用输出2,后者使用输出1,它们的Mean Dice, F1值和Hausdorff距离并无较大改变,基本保持在同一水平线,但Sub的参数量减少了1 M左右,在保证分割精度的同时减少测试时参数。
表2 消融实验
(3)不同注意力模块对比实验。为验证本文提出注意力模块的有效性,分别将BAB中的注意力模块舍弃或替换为scSE进行对比实验。输出采用输出1,实验结果如表3所示。对比本文方法与不添加注意力机制、添加scSE的分割结果,本文方法的提升效果最为明显。相较于scSE,在各数据集上Mean Dice提升0.3%左右,F1值提升1%左右,Hausdorff距离也都有所减小。可见本文设计注意力模块的有效性,可强化学习重点特征,弥补缺失值,融入在BAB中,进一步提高网络分割精度。
表3 不同注意力模块在不同数据集上的分割结果对比
(4)不同损失函数对比实验。为了验证本文提出组合损失函数的有效性,在AS-UNet网络中分别对基于区域的Dice损失、基于边缘的Boundary损失和本文所提Dice + Boundary 损失这3者进行对比实验。实验采用Mean Dice, F1值和Hausdorff距离作为评价指标。3个数据集上的实验结果如表4所示。从表4中结果可见,组合损失函数在各数据集上都能达到更好的分割效果,相较于单一的Dice损失和Boundary损失,Dice + Boundary损失的Mean Dice和F1值均提高了1%左右,Hausdorff距离均减小2左右。这是因为本文方法引入掩膜边缘图,增强边缘信息,结合区域以及边缘的损失函数能够更真实地反映预测图与GT之间的差异,从而进行参数校准与更新。
表4 不同损失函数在不同数据集上的分割结果对比
4 结束语
本文针对医学图像分割缺失值大、边缘不清晰、模型参数冗余等问题提出AS-UNet。训练时引入边缘注意模块,将掩膜边缘图连接到UNet扩张路径的最后3层上以补充边缘信息,解决边缘分割不清晰的问题;同时在边缘注意模块中引入新的注意力模块,并结合多尺度特征图以强化重要特征学习,弥补部分缺失值;提出了新的组合损失函数优化网络参数,并实现测试时只需要选择训练好模型的输出1即可得到最终结果,在提高分割精度的同时减少测试时参数。实验结果表明,本文模型在Glas, DRIVE, ISIC2018 3个不同类型的分割数据集上都有较好表现,验证了本文方法的有效性和泛化能力。论文相关代码地址:https://github.com/gqq1210/AS-UNet。