改进U-Net的多级边缘增强医学图像分割网络

胡帅,李华玲,郝德琛

(中北大学软件学院,山西 太原 030051)

0 引言

随着现代医学的不断发展,各种医疗影像设备广泛应用于临床诊疗,计算机断层扫描(CT)、X射线、超声、正电子发射断层扫描(PET)和核磁共振成像(MRI)等均为医师诊断治疗发挥着重要的作用。医学图像分割的主要目的是从各类医学图像中自动分割出病灶区域,辅助医师更快更准确地作出诊断,极大地提高了医师的诊断效率。因此,医学图像分割模型的分割效果显得格外重要。与此同时,医学图像相对于自然图像背景环境更复杂,具体表现在结构边缘模糊、形状多变、纹理结构细节丰富,使得医学图像分割任务更具挑战性。随着深度学习技术在医学图像分割中的广泛应用,医学图像分割精度有了显著提升。卷积神经网络(CNN)的出现,尤其是全卷积神经网络的出现,使深度学习模型具备了高速度、高精度和全自动的特点。RONNEBERGER等[1]提出用于医学图像分割的卷积网络(U-Net)。鉴于U-Net的编解码器结构在医学图像分割中的优异效果,多数性能优异的医学图像分割模型都是基于U-Net的主要结构进行改进,例如:用于医学图像分割的嵌套U-Net(UNet++)[2],注意力U-Net(Attention U-Net)[3],用于医学图像分割的全尺寸连接UNet(UNet 3+)[4],基于多层感知器的快速医学图像分割网络(UNeXt)[5]等模型。

目前,医学图像分割模型仍存在以下问题:1)普通卷积运算的局部性导致多数模型未能有效地提取图像的全局信息;2)对图像边缘信息的忽略导致病灶分割边缘区域模糊。为解决以上问题,本文在保留U-Net的高效编解码器结构的同时,在U形结构跳跃连接部分引入多级边缘增强(MEE)模块和细节特征关联(DFA)模块,提出改进U-Net的多级边缘增强医学图像分割网络(MDU-Net)模型,改善了医学图像分割时边缘模糊的问题,同时提高了分割精度,在LiST2017肝脏肿瘤分割挑战数据集[6]、胸部X射线掩膜和标签数据集[7-8]2个数据集上进行实验验证。

1 相关工作

目前,主流的医学影像分割模型主要分为CNN和Transformer 2种结构。CNN自提出以来已成为图像处理的主流结构。2015年,SHELHAMER等[9]提出全卷积网络(FCN),开启了其在语义分割领域的应用,FCN的基本思想是使用卷积层替换全连接层,使用转置卷积进行上采样,分割精度优于传统的分割方法。在FCN的基础上,RONNEBERGER等[1]通过引入跳跃连接提出一种编解码器结构的U-Net,编解码器之间通过跳跃连接将相同层级的特征信息进行拼接,U-Net在各种类型的医学图像分割任务中均表现出了十分优异的分割性能,可以有效地融合低层和高层图像特征,缓解语义间隙问题,在医学图像上显示出很好的检测能力。由于U-Net在医学图像分割领域的突出贡献,随即出现了多种改进U-Net的医学图像分割网络。OKTAY等[3]提出Attention U-Net,能够关注不同形状和大小的目标结构,同时能够抑制图像中的不相关区域并突出有用的显著特征。ALOM等[10]提出一种用于医学图像分割的递归残差U-Net(R2U-Net),使用递归卷积代替传统卷积,以保证更好的特征表示。林志洁等[11]使用内卷操作替换了原始的卷积操作提升了U-Net对局部特征的学习能力。刘文等[12]对U-Net架构中的卷积模块进行改进,同时引入了批量归一化层提出Concat-UNet,增强了网络的特征提取能力。LIN等[13]提出一种改进FCN的网络架构(RefineU-Net),通过结合浅层特征中的全局上下文信息和深层特征中的语义信息,将其进行全局细化,使用生成的判别性注意力特征与解码器特征进行连接,便于局部细化操作。HUANG等[14]提出一种密集卷积网络(DenseNet),缓解了梯度消失问题,加强了特征传播,提高了特征的重复利用。受此思想的启发,AZAD等[15]使用双向卷积长短期记忆(BConvLSTM)网络并以非线性的方式将从相应层次的编码器与解码器中提取的特征图结合起来,采用密集连接的卷积操作实现特征重用,提出一种用于医学图像分割的双向网络(ConvLSTM U-Net),优化跳跃连接,同时结合BConvLSTM块,以捕获更有鉴别性的信息。ZHOU等[2]提出UNet++,以缩小编解码器之间的特征图语义差距。HUANG等[4]提出UNet 3+,将全尺寸跳跃连接与不同尺度的特征信息图的高级语义相结合,并通过深度监督从全尺寸聚合特征图中学习分层表示。HE等[16]提出残差结构,增加了从输入到输出的路径,有效缓解了梯度消失等问题,提高了网络的收敛速度和运行效率。为了解决卷积块堆叠造成的有效感受野减小的问题,XU等[17]通过在普通卷积块之前嵌入由深度可分离卷积组成并具有残差结构的卷积模块,提出一种用于医学图像分割的更深入紧凑的分散注意力U-Net(DCSAU-Net)。

在医学图像分割领域,自动化的分割模型不仅对分割的准确率有较高的要求,而且对于边缘区域的分割也要求严格。对于医学图像分割边缘区域模糊的问题,学者们进行了大量研究。LIU等[18]结合边缘特征提取模块提出一种用于医学图像分割的多层边缘注意力网络(MEA-Net),通过从不同的编码器阶段采用卷积操作获取的边缘特征序列融合到最后的分割掩码中,提高边缘区域的分割效果。ZENG等[19]提出一种用于医学图像分割的密集边缘注意力UNet(DEA-UNet),该模型从U形网络最底部上采样层开始使用反向注意力模块逐级提取特征,并结合边缘引导模块更大限度地获取边缘特征。ZHANG等[20]在编码器部分加入了边缘注意力引导模块来获取边缘信息表示,最后将边缘特征与多尺度编码器的输出特征信息相结合生成最后的分割掩码。HAO等[21]提出一种基于图的边缘注意力门医学图像分割模型(EAGC-UNet++),将图论的思想引入医学图像分割任务中,使用残差图卷积块以及传统卷积块共同作为UNet++网络的编码器部分,在模型中加入边缘注意力门来获取边缘信息。

在上文提到的各种改进U-Net模型中,对于医学图像的边缘信息处理还有待改进,医学图像的细节特征信息利用还不够完善。为了解决上述问题,本文提出一种MDU-Net,由扩张门控注意力(DGA)、多级边缘增强、细节特征关联3个新设计的模块组成。

2 MDU-Net模型

2.1 整体架构

本文提出的MDU-Net模型的整体架构如图1所示。以编解码器结构的U-Net模型作为主体结构,并嵌入MEE模块和DFA模块。多级边缘增强模块的设计是为了充分捕获低级特征图中的边缘特征信息,融合不同层次的图像边界特征,为最终的分割效果提供边界补充信息。细节特征关联模块接收双层特征图信息,逐步引导多尺度上下文特征信息的融合,高效利用图像的深层次细节特征。在主体结构中加入各个模块以学习图像的边界以及细节特征,能获取更准确的分割结果。

图1 MDU-Net模型整体架构Fig.1 Overall architecture of MDU-Net model

2.2 多级边缘增强模块

在编码器结构中,低层特征具有丰富的空间细节特征,高层特征具有丰富的语义信息。孙军梅等[22]仅在U-Net基准网络的解码器的最后3层进行了边缘特征的提取,忽略了编码器低级特征层中丰富的边缘特征信息,会导致大量边缘信息缺失。不同层次特征信息的融合不仅实现了多层次特征间的优势互补,而且有效地提高了网络的远程上下文学习。在U-Net网络中,在编码器结构中的下采样操作会造成边缘特征信息的严重缺失。为了补充缺失的边缘信息和高效利用低级特征中丰富的边缘信息,提出MEE模块,该模块对低级特征的边缘信息进行深入提取并补充至解码器特征信息中。如图2所示,该模块将编码器结构中低级特征块E1输入扩张门控注意力模块获得高级边缘信息块E′1,将特征块E2先输入改进的注意力融合(AF)模块在获得特征信息后进行一次上采样操作得到E′2,再与高级边缘信息块E′1进行元素相乘得到边缘增强信息Ed。

图2 多级边缘增强模块结构Fig.2 Structure of multistage edge-enhanced module

上述过程如式(1)～式(3)所示:

E′1=DGA(E1)

(1)

E′2=UpSampling(AF(E2))

(2)

Ed=E′1⊗E′2

(3)

其中:DGA为输入扩张门控注意力模块;UpSampling表示上采样操作;⊗表示基于元素的乘法。

在改进的并行注意力模块中,将通道注意力块和空间注意力块进行并联操作,通过参数α和β控制注意力块的加和比例,参数α、β的总和为1。经过多次实验验证,设置α、β同时为0.5可获得最佳分割效果。

2.3 扩张门控注意力模块

医学图像分割属于密集预测任务,同时获取全局和局部信息是提高分割性能的关键。全局信息有助于模型了解目标区域的整体结构及其与背景的关系,从而更准确地定位目标区域。李翠云等[23]将边缘增强的注意力模块嵌入解码器结构,未注意到低级特征中包含的边缘信息。本文注意到低级特征中的丰富边缘信息,同时采用不同扩张率的深度可分离卷积获取不同的特征信息。首先,通过扩张率为1和3的卷积获得局部特征信息,通过扩张率为5和7的卷积获得全局特征信息。然后,将取得的局部特征信息和全局特征信息进行逐像素相加,并加入残差信息。这样能有效建立图像的上下文信息,充分利用不同维度的特征信息对全局信息进行补充。

上述过程可由式(4)～式(7)表示:

Edi=DSCd=i(E)

(4)

Eattg=Sigmoid(Ed7)⊙Ed5

(5)

Eattl=Sigmoid(Ed3)⊙Ed1

(6)

Edga=E+Eattl+Eattg

(7)

其中:E表示输入特征;DSCd=i(i=1,3,5,7)表示扩张率为i的深度可分离卷积操作;⊙表示基于元素的乘法;Eattg和Eattl为计算出的全局特征注意力图和局部特征注意力图。

扩张门控注意力模块结构如图3所示。

图3 扩张门控注意力模块结构Fig.3 Structure of dilated gated attention module

2.4 细节特征关联模块

医学图像分割目标的尺度有很大变化,同时一些分割目标结构边界模糊,形状复杂,纹理异质,非常容易产生错误的分割结果。多尺度上下文信息可以提高语义分割任务的性能,由于目标的尺度是多样的,因此有效利用多尺度的上下文信息是实现目标精确定位的关键。受此启发,在细节特征关联模块中,嵌入多尺度特征提取(MFE)模块来获取特征图的多尺度特征信息,并联平均池化和最大池化操作,将池化结果相加后进行Sigmoid操作生成池化特征注意力图Epoolatt,如图4所示。

图4 细节特征关联模块Fig.4 Detail feature association module

WOO等[24]认为:通道注意力和空间注意力在抑制不相关信息方面发挥着重要作用。王晓援等[25]在高级特征信息块中提取了多尺度特征信息,但未考虑模型对高级特征信息的注意程度,会提取过多冗余的信息。本文在特征注意力融合模块中并联通道注意力模块以及空间注意力模块最终获取注意力融合特征图,以此提高该模块对相关信息的注意程度,最后将池化特征注意力图与MFE模块得到的多尺度特征信息进行逐像素相乘,之后与注意力融合特征图进行逐像素相加获得该模块的结果。

上述过程可由式(8)～式(14)表示:

Epoolatt=

Sigmoid(AvgPool(Ei)+MaxPool(Ei))

(8)

M1=Conv1(Ei+1)

(9)

M3=Conv1(Conv3(Ei+1))

(10)

M5=Conv1(Conv5(Ei+1))

(11)

M7=Conv1(Conv7(Ei+1))

(12)

Esc=Concat(M1,M3,M5,M7)

(13)

Eout=Epoolatt⊙Esc+AF(Ei+1)

(14)

其中:Ei表示输入特征;AvgPool和MaxPool分别表示平均池化操作和最大池化操作;Concat表示特征图通道连接操作;⊙表示基于元素的乘法;AF表示注意力融合操作;Eout表示该模块最后的输出结果。

3 实验与结果分析

3.1 数据集与数据预处理

使用2个数据集来验证MDU-Net模型的通用性和有效性。由于获取图像的方式大相径庭,因此在使用同一模型进行训练时需要对原有的图像进行预处理,将图像的长和宽处理为相同尺寸。在使用数据时,将图像统一设置为256×256×3,即图像长度和宽度分别为256像素,通道数为3。

3.1.1 LiTS2017肝脏肿瘤分割挑战数据集

LiTS2017肝脏肿瘤分割挑战数据集包含131组由世界各地医院提供的CT图像数据,选择其中前30组CT图像数据进行实验。在实验中,将数据集按80∶20划分训练集和测试集,其中,训练集包含1 540张图像,测试集包含385张图像。

3.1.2 胸部X射线掩膜和标签数据集

胸部X射线掩膜和标签数据集包含138张X射线图片,其中,80张X射线图片正常,58张X射线图片存在结核病异常。这个数据集有丢失标签的问题,经过分析和处理后,将数据集按80∶20划分训练集和测试集,其中,训练集包含396张图像,测试集包含99张图像。

3.2 实验处理

3.2.1 实验环境

模型使用Python 3.8和PyTorch实现。模型训练和测试是在1台服务器上进行的,使用NVIDIA RTX 3060Ti显卡和16 GB内存。将数据集中的所有图像及其对应的标签大小分别调整为256×256×3。使用的优化器是RMSProp,学习率为0.000 1。用于训练的批量大小为2,轮次数设置为50。

3.2.2 评价指标

为了评估MDU-Net模型的有效性,采用Dice相似系数、灵敏度、准确率、特异性、查准率以及Jaccard相似系数作为评价指标,计算公式分别如式(15)～式(20)所示:

(15)

(16)

(17)

(18)

(19)

(20)

其中:NTP、NTN、NFP、NFN分别表示真阳性、真阴性、假阳性、假阴性样本数量。

上述评价指标的数值越大,代表模型分割效果越好。如果这些指标的数值接近1,则意味着输出掩码与标签真值掩码重叠。

3.3 实验结果

3.3.1 模型比较

1)LiTS2017肝脏肿瘤分割挑战数据集。分别使用U-Net[1]、UNet++[2]、DCSAU-Net[17]、利用深度残差U-Net的道路提取(ResUNet)模型[26]、用于二维医学图像分割的上下文编码网络(CE-Net)模型[27]、用于图像分割的深度卷积编解码器体系结构(SegNet)[28]、用于医学图像分割的Transformers强编码器(TransUNet)模型[29]、用于医学图像分割的Transformer编码器和CNN解码器的整合(LeViT-UNet)模型[30]等在LiTS2017数据集上进行实验,实验结果的数据展示如表1所示,其中最优指标值用加粗字体标示,下同,对应的分割结果如图5所示。由表1可以看出,MDU-Net实现了最好的效果,在准确率、灵敏度、Jaccard相似系数、Dice相似系数等指标上均名列首位,尤其在灵敏度、Dice相似系数指标上有明显的提升,分别达到了96.24%、0.926 5,在准确率、Jaccard相似系数指标上均有提升。由图5可以看出,MDU-Net对目标区域的分割效果在完整度以及边缘分割细节上都明显优于其他模型,再次证明了MDU-Net相较于其他模型更加高效。

表1 LiTS2017肝脏肿瘤分割挑战数据集上的实验数据Table 1 Experimental data on the LiTS2017 liver tumor segmentation challenge dataset

图5 LiTS2017肝脏肿瘤分割挑战数据集上的分割结果比较Fig.5 Comparison of segmentation results on the LiTS2017 liver tumor segmentation challenge dataset

2)胸部X射线掩膜和标签数据集。使用不同模型在胸部X射线掩膜和标签数据集上进行实验,最终的实验数据和实验分割效果分别在表2和图6中进行展示。由表2可以看到,MDU-Net在准确率、特异性、查准率、Jaccard相似系数、Dice相似系数评价指标上都有提升,分别达到了97.96%、98.53%、95.46%、0.920 4、0.957 8。由图6可以看出,MDU-Net在分割目标边缘区域具有优秀的分割效果,主要原因为MEE模块提供了更多的边缘信息,使得MDU-Net在边缘处的效果显著优于其他模型。

表2 胸部X射线掩膜和标签数据集上的实验数据Table 2 Experimental data on chest X-ray mask and label dataset

图6 胸部X射线掩模和标签数据集上的分割结果比较Fig.6 Comparison of segmentation results on the chest X-ray mask and label dataset

此外,还对MDU-Net的参数量、计算量指标与不同模型进行比较分析,结果如表3所示。MDU-Net为了提取医学图像中更为复杂丰富的特征信息,提出的MEE、DGA和DFA模块结构相对复杂,导致模型的参数量和计算量相对较高,但这是一个值得权衡的问题。在医学图像分割领域,分割精度对于诊断和治疗意义重大,因此提高分割精度是至关重要的目标。结合图5和图6的分割结果可以看出,增加的参数对于模型的分割效果起到了关键作用。

表3 模型参数量以及计算量分析比较Table 3 Analysis and comparison of the model parameter and calculation quantity

3.3.2 消融实验

为了进一步验证所提模块的有效性,在实验条件相同的情况下,以U-Net为基准网络,逐步添加多级边缘增强模块、扩张门控注意力模块和细节特征关联模块进行消融研究,并在肝脏肿瘤分割挑战数据集和胸部X射线掩模和标签数据集上分别进行实验,实验数据如表4和表5所示,其中“√”表示使用该模型或模块。

表4 LiTS2017肝脏肿瘤分割挑战数据集上的消融实验分割结果Table 4 Segmentation results of ablation experiments on the LiTS2017 liver tumor segmentation challenge dataset

分析两组消融数据可知,相比于基准网络U-Net,在其添加DFA模块后的准确率、灵敏度、Jaccard相似系数、Dice相似系数指标上均有所提升。DFA、DGA和MEE 3种模块的组合在整体上的分割效果最好,基于3种模块的网络模型在胸部X射线掩模和标签数据集的准确率、特异性、查准率、Jaccard相似系数、Dice相似系数上表现均为最佳。

4 结束语

本文注重高效保留医学图像的边缘信息以及细节特征的提取与融合,提出改进U-Net的MEE医学图像分割网络模型。在具有丰富边缘特征信息的低级特征块中,利用MEE模块提取边缘特征,嵌入DGA模块用于抑制非重要特征的获取。在此基础上,在模型跳跃连接上嵌入DFA模块用于补充图像的全局和局部特征信息。实验结果表明,该模型在边缘区域分割以及整体分割效果上表现优异,优于目前先进的医学图像分割模型。下一步将继续优化所提模型,降低模型参数量,使其可应用于实际医学图像分割任务。