PRA-UNet3+:全尺度跳跃连接CT肝脏图像分割模型
2023-03-08钟经纬
钟经纬
(江南大学 人工智能与计算机学院,江苏 无锡 214122)
0 引言
近年来,由于不健康的饮食与作息习惯,全球肝癌发病率持续上升,据世界卫生组织/国际癌症研究机构统计报告[1],2020 年全球肝癌新发病例数为90.6 万例,其中中国占全球新发肝癌病例的45.3%,肝癌已经成为影响我国居民健康的重大公共卫生问题。多种原因都可导致肝癌发病率升高,且肝癌起病隐匿、早期发现困难、中晚期治疗复杂且效果差异大,因此当今肝癌治疗依然是一个比较困难的医学问题。随着医学影像技术与计算机技术的发展,许多研究者尝试用计算机技术解决癌症患者的早期确诊问题[2-3],为癌症治疗争取宝贵时间。器官定位与图像分割是诊断治疗的第一步,但由于医学图像,尤其是CT 图像对比度较低,且目标不明确,医生在诊断时通常需耗费大量时间与精力来估计器官的位置和大小,个人经验也会影响诊断结果,导致在严肃的医学诊断结果中夹杂了主观性。因此,应用计算机技术辅助诊断,构建一种精确、快速的器官图像分割方法对癌症诊断及早期治疗具有重要意义。
1 相关工作
当前的器官图像分割方法很多,主要分为传统图像分割方法和基于深度学习的图像分割方法。
1.1 传统图像分割方法
传统图像分割方法主要包括基于阈值、边缘和图论的分割方法。基于阈值的分割方法的基本思想是基于图像灰度特征计算一个或多个灰度阈值,并将图像中每个像素的灰度值与阈值相比较,最后将像素根据比较结果分到合适的类别中。因此,该类方法最为关键的一步就是按照某个准则函数求解最佳灰度阈值。Prewitt 等[4]提出一种全局单阈值分割方法,即直方图双峰法(mode 法)。基于边缘的分割方法基于灰度值的边缘检测,是建立在边缘灰度值会呈现出阶跃型或屋顶型变化这一观测基础上的方法。阶跃型边缘两边像素点灰度值存在明显差异,而屋顶型边缘位于灰度值上升或下降的转折处。Canny[5]提出一种边缘角点和兴趣点检测器,命名为Canny 边缘检测器。基于图论的分割方法将图像分割问题与图的最小割问题相关联,其本质就是移除特定的边,将图划分为若干子图,从而实现分割。目前基于图论的方法有GraphCut[6]、Grab-Cut[7]、Random Walk[8]等。
1.2 基于深度学习的图像分割方法
近年来,随着计算机技术的快速发展,深度学习技术开始大量应用于医学图像分割领域。早期,半自动分割[9]是一种常见的分割手段。在进行医学图像分割时,因为感兴趣目标的不同以及不同类型图像的差异,没有一个普遍适用的分割方法。对于同一模态的图像,如果不能充分利用感兴趣目标和一些先验信息,仍然无法获得精度较高的分割结果。该阶段的分割手段往往需要一定的手工干预,才能得到令人满意的分割结果。随后,基于卷积神经网络(Convolutional Neural Network,CNN)或循环神经网络(Recurrent Neural Network,RNN)的分割方法开始涌现,可以快速、准确地获取图像的自然分割结果。Long 等[10]提出全卷积神经网络(Full Convolutional Neural Network,FCN)用于图像语义分割,并取得了很好的效果。FCN 能够准确、完整地定位与分割医学图像中的器官和病变部位,对患者的后续治疗具有重要意义,因此得到了迅速发展。Olaf 等[11]提出的U-Net 网络结构特别适用于医学图像分割领域,采用编码解码器结构,通过跳跃连接巧妙地结合了高层和低层信息,解码层的深层抽象信息能更好地利用编码层传输的浅层信息,使得图像分割效果较好。
随后,许多学者对FCN 和U-Net 进行了研究,并在此基础上提出许多新的用于医学图像分割各细分领域的网络结构。例如,Bi 等[12]提出一种新的ResNet 级联网络,融合多尺度信息逐步精确定位肝脏图像边界;DenseNet[13]改进了卷积块之间的连接方式,使每个卷积块都与其之后的卷积块残差相连;Kaluva 等[14]提出一种全自动的两阶段级联网络,第一阶段进行肝脏图像分割,第二阶段使用第一阶段的分割结果进行肝脏定位;Li 等[15]提出H-DenseUNet探索肝脏和肿瘤图像分割的混合特征;Christian 等[16]提出一种改进的CT 和MRI 心脏图像分割方法;Liao 等[17]提出一种利用卷积神经网络检测CT 图像中恶性结节的方法。另外,一些通用的医学图像分割网络也相继被提出,例如Alom 等[18]提出R2U-Net,将剩余单元与RNN 结合;Diakogiannis 等[19]提出ResUNet-a,该网络是对R2U-Net 的一种改进,增加了剩余连接数,并结合了金字塔场景分辨率池和多任务推理;Zhou 等[20]提出UNet++,主要对跳转连接部分进行了改进;Li等[21]引入注意记忆,并提出了ANU-Net;Huang 等[22]通过结合多尺度能力改进U-Net,提出了UNet 3+网络模型。近年来基于Transformer 的医学图像分割方法也不断涌现,例如Chen 等[23]提出TransUNet,首次将Transformer 引入医学图像分割领域,一方面把编码器卷积块替换为Transformer 块,提取全局上下文的输入序列信息,另一方面基于卷积的解码器进行上采样,然后与高分辨率的卷积特征图相结合,实现精确定位;Valanarasu等[24]提出Medical Transformer,采用局部—全局训练策略(LoGo),进一步提升模型分割性能。
为了提高医学图像分割精度,本文在 UNet 3+网络基础上引入注意力门和点采样技术。注意力门可使网络在训练过程中更加关注分割目标区域,点采样方法可提高模型判断分割目标边缘像素类别的能力,改善目标边缘分割效果。通过在肝脏医学图像数据集CHAOS[25]上进行实验,本文提出的模型PRA-UNet3+都表现出一定的性能优势,在视觉感知上的分割效果也有明显提升。
2 网络结构
2.1 注意力门
注意力门(Attention Gate)是由Oktay 等[26]提出的,在网络中引入注意力门,可帮助模型增强对目标区域的学习。注意力门通过同一级编码器特征与深层解码器上采样特征,学习到一个关注分割目标结构的注意力系数,以帮助分割网络的重点关注位置。因为不断进行下采样,获得了更大的感受野,但细节被抽取,损失了语义信息,而浅层编码器的信息更加丰富,所以浅层的编码器信息对注意力系数的学习也是有帮助的。本文注意力门的结构设计不仅融合了同一层编码器信息,而且进一步融合了所有比该层更浅层的编码器信息。注意力模块如图1所示。
Fig.1 Attention module图1 注意力模块
PRA-UNet3+中注意力门计算过程如下:①注意力门有两个输入:一是编码器特性(f),该特性由浅层编码器和同层编码器特征组成,二是上采样信号(g),该信号是所有深层解码器上采样信息特征融合的结果;②经过卷积运算(Wf,Wg)和BatchNorm 运算(Bf,Bg)后,将其各自的输出进行特征融合;③特征融合结果随后被输入到激活函数(ReLU,σ1(x)=max(0,x))中;④被激活后,该特征再次经过卷积运算(Wθ)和BatchNorm 运算(bθ);⑤将结果输入激活函数计算注意系数(α),让参数在注意力门迅速收敛;⑥最后将注意力系数(α)逐点乘以编码器特征以获得输出。
式(1)-式(3)总结了注意力门计算过程:
注意力门可获得全局与局部的联系,增强重要信息的学习,并抑制不相关信息的学习。此外,与CNN 和RNN 相比[27],其参数更少,降低了模型的复杂性。
2.2 点采样方法
在微观层面上,图像的低层由像素组成,每个像素可用一组特征向量表示。利用线性插值等方法,将特征向量均匀映射到一组标签上,称为上采样。相反,本研究使用经典的细分策略[28],在该策略中不是对所有像素进行上采样,而是根据特定策略选择一定数量的不确定像素进行特殊处理,该方法被称为点采样。该方法对分割边缘比较难预测准确的点进行二次细化预测,可提高分割边缘点的判断正确率。
医学图像分割同样存在边缘分割困难的问题,对边缘进行单独细化处理,可解决模糊分割边界问题,提高分割性能及视觉效果。具体实施步骤如下:①当前解码器从深层解码,输入粗掩模特征图;②点采样模块依据点选择策略从粗掩模特征中选择最不确定的一批采样点;③用一个多层感知机(Multilayer Perceptron,MLP)对最不确定的采样点进行细分预测,ReLU 用于激活MLP 的隐藏层,Sigmoid用于激活输出;④点特征映射到编码器特征(虚线箭头),并替换相应位置的特征,以获得点采样特征;⑤将上采样特征、注意力门输出特征与点采样特征进行特征融合,得到细掩模特征图。
点采样模块计算过程如图2 所示。采样点主要在分割目标边缘附近,且具有一定的覆盖性,这样才能达到较好的点采样效果。传统点选择策略是从均匀分布的像素点中随机抽取kN(k>1)个点,但由于医学图像数据正负样本不平,如背景面积大、分割目标面积小,这种点选择策略容易造成过采样。因此,本文提出根据正负样本点比例不同进行不同比例的点采样。
Fig.2 Point sampling module图2 点采样模块
点采样策略如下:①正样本数量Np,负样本数量Nt,求出正负样本比例;②按照图像分辨率大小设置一个采样常数N,采样点个数为mN,当正样本数量多时,边采样点多,反之边采样点少;③对上述mN个采样点进行粗略预测和插值操作,计算所有点的不确定度,重点关注最不确定的点,选出amN个最不确定的点;④除去最不确定的点,剩下(1 -a)mN个不重要的点并对其进行简单的上采样操作;⑤最后获得边缘精确分割的掩模。
2.3 网络总体结构
U-Net网络框架及其变体在医学图像分割任务中得到了广泛应用,并取得了很好的效果。其中,U-Net++和ANU-Net 是两种代表性的体系结构。Huang 等[22]在UNet++基础上,为将不同尺度特征图的浅层语义信息与高层语义信息相结合,提出U-Net 3+网络。PRA-UNet3+在UNet3+模型基础上进行了修改,添加注意力门模块和点采样模块。PRA-UNet3+网络架构主干如图3所示。
Fig.3 PRA-UNet3+network structure图3 PRA-UNet3+网络结构
U-Net、U-Net++及其变体无法获取图像的全尺度信息,因此无法学习分割目标的位置,分割边界相对模糊。为了弥补这一不足,PRA-UNet3+借鉴U-Net 3+结构,网络中每个解码器都结合了所有浅层和深层编码器的输出特征,以及深一层编码器的点采样结果。组合这些信息可以捕获细粒度和粗粒度的语义信息,但是浅层信息与深层信息以何种比例组合才能更好地关注到分割目标区域仍是一个问题。因此,本文提出在浅层与深层特征间加入一个注意力门来学习一个注意力系数,决定浅层信息与深层信息特征融合的比例,更好地关注分割目标区域。网络结构中使用的点采样模块为改进的点采样模块(IPR)。
如图4 所示,以获取XDe_3的特征图信息为例,阐述网络组成细节。与U-Net++一样,接受同尺度编码器层的特征图XEe_3,该特征需经过64 个大小为3 × 3 的滤波器进行卷积运算,目的是统一特征图数量,减少冗余信息。浅层编码器XEe_1和XEe_2通过最大池化进行下采样,目的是统一特征图分辨率。从图4 中可知,XEe_1缩小为原,XEe_2缩小为原,随后同样经过滤波器进行卷积运算。经过运算后的XEe_1、XEe_2和XEe_3进行特征融合,即为注意力门编码器特征输入端。深层解码器XDe_4和XDe_5利用双线性插值进行上采样,分辨率分别放大2 倍和4 倍,同样需要经过64 个大小为3 × 3 的滤波器进行卷积运算来统一特征图数量。将上述经过运算后的XDe_4和XDe_5进行特征融合,即为注意力门解码器特征输入端。注意力门模块获得两端输入后进行模块内部运算操作,步骤如2.1 节所述。将原XDe_4输入到改进的点采样模块进行计算,获得细掩模特征,步骤如2.2 节所述。将注意力门模块输出特征和点采样模块输出特征进行特征融合,经过320 个大小为3×3 的滤波器进行卷积运算、批量归一化和ReLU 激活函数激活便获得了XDe_3。
Fig.4 PRA-UNet 3+decoder calculation图4 PRA-UNet 3+解码器计算
3 实验结果及分析
实验平台为Intel Core i9-9900X 处理器,NVIDIA Ge-Force RTX 2080Ti 显卡,Ubuntu 16.04操作系统。
3.1 数据集与预处理
本文使用CHAOS作为实验数据集,来自2019 ISBI调整赛。数据集包含CT和MRI两种扫描数据,本文仅采用CT数据作为实验数据,包含40 个不同病人的腹部CT 扫描图像。拆分后汇总得到2314个有标签样本,分辨率为512 × 512。
为了节省计算资源,便于训练模型,把图像分辨率调整为128 × 128,将训练集、验证集和测试集按照6:2:2 的比例进行划分。实验前对实验数据进行常规的数据增强,例如随机翻转、平移缩放等,以缓解小数据集带来的过拟合现象。
3.2 评价指标
为评估模型性能以及与其他模型进行比较,本文采用Dice 系数(Dice)、交并比(IoU)和 F1 分数(F1 Score)作为评价指标。F1 Score 指标可同时体现精确率(Precision)和召回率(Recall),通常用来评价二分类模型的稳健程度。在式(4)、式(5)中,X为预测结果,Y为真实标签值。在式(6)中,TP为真阳性,FP为假阳性,FN为假阴性。这3 个指标的值越大,与实际结果的相似度越高,则分割结果越好。
3.3 实验结果及分析
为验证本文提出的医学图像分割模型在肝脏图像分割数据集CHAOS 上的性能优势,将其与其他几种经典医学图像分割模型进行比较,包括U-Net[11]、U-Net++[20]、Attention U-Net(AU-Net)[21]、Medical Transformer(MedT)以及U-Net 3+[22]。表1 为不同网络模型在肝脏图像分割数据集上的性能比较。
Table 1 Comparison of liver image segmentation performance表1 肝脏图像分割性能比较
从表1 可以看出,本文提出的医学图像分割模型PRA-UNet3+在肝脏图像分割任务上,3 项评价指标均优于其他模型。在测试集上平均Dice 达到0.9467,平均IoU 达到0.9623,平均F1 Score达到0.9351。
由表1 可知,U-Net 3+融合全尺度的语义信息,在性能表现上优于U-Net、U-Net++、AU-Net 这类普通的跳跃连接架构,证明了全尺度跳跃连接架构的有效性。由于浅层语义信息中冗余信息较多,因此之前的模型在分割目标边缘处理以及分割目标与背景对比度低时,分割效果并不理想。本文在U-Net 3+浅层语义信息与深层语义信息融合时,加入注意力门来过滤冗余信息的影响,并引入点采样方法以解决传统采样步骤对分割目标边缘过采样的问题,使得边缘分割更加平滑。对比U-Net 3+,PRA-UNet3+平均Dice 提升了2.55%,平均IoU 提升了2.27%,平均 F1 Score 提升了2.68%,证明该模型能更好地对肝脏图像和分割背景进行分类,提高分割精度。为了更加直观地对比分割效果,图5 给出4 组肝脏图像在不同模型下的分割结果比较。
Fig.5 Comparison of liver image segmentation results图5 肝脏图像分割结果比较
方框标出分割效果差异较大的部分,从图5 中可以看出PRA-UNet3+对分割边缘的处理上更加平滑,分割正确的面积更大,避免了过分割和欠分割问题,与真实标签值最为接近。
为证明模块的有效性,同时进行了消融实验。以UNet 3+作为基线网络,依次将注意力门模块(AG)和改进的点采样模块(IPR)融入其中进行实验。
(1)模块有效性。表2 中的数据显示,在CHAOS 数据集上,当U-Net 3+只加入浅层与深层语义信息之间的注意力模块时,Dice 为0.9174,相比U-Net 3+提升了6.27%;IoU 为0.9234,相比U-Net 3+提升了3.61%;F1 Score 为0.9140,相比U-Net 3+提升了 6.33%。可见引入注意力模块的网络在各个评价指标上均有明显提升,说明通过注意力门模块可以提升网络对重点分割区域的关注程度。若去除注意力模块,只在解码过程中加入改进的点采样模块,使网络对边缘难分割的区域进行细化分割,此时Dice、IoU 和F1 Score 的值分别为0.9232、0.9449 和0.9187,各指标分别比U-Net3+网络提升了6.94%、6.03%和6.88%。因此,单独引入注意力模块或改进的点采样模块,对分割性能都有一定提升,证明了各模块的有效性。详细的消融实验结果比较如表2所示。
Table 2 Results of ablation experiment on CHAOS表2 CHAOS数据集上的消融实验结果
(2)改进的点采样模块有效性。由表2 可知,相比传统点采样模块(PR),改进的点采样模块在一定程度上缓解了类样本不平衡问题,3 个评价指标都有所提升,Dice 提升较少,仅为0.87%,IoU 提升了2.68%,F1 Score 提升了1.08%。由此证明,改进的点采样模块在本文提出架构中是有效的。
模型大小关系到实践的可能性,也是评价一个模型的重要指标。利用PyTorch 工具包对本文模型与其他模型在参数量方面进行比较,结果如表3所示。
Table 3 Comparison of model parameters表3 模型参数量比较
表3 给出不同模型参数量对比。U-Net 3+由于使用全尺度的跳跃连接,参数量明显少于 U-Net。本文模型基于U-Net 3+,参数在注意力模块和点采样模块有所增加,参数量相比U-Net 3+增加了1.49%,但Dice 提高了2.55%,IoU 提高了2.27%,F1 Score 提高了2.68%,以较小的计算性能消耗为代价,带来了比较可观的性能提升,证明了本文模型在肝脏图像分割任务上的优越性。
4 结语
器官图像分割是医学图像技术对病理分析的基础,也是进行下一步深入研究的前提,所以一个好的分割方法对医学图像研究至关重要。本文在经典U-Net 3+分割模型上进行改进,在全尺度特征融合过程中加入注意力门,使模型更加关注分割目标区域。为了解决传统采样方法造成目标分割边缘过采样的问题,在上采样过程中对分割边缘进行点采样处理,从而改善边缘分割判断不准确的现象,提高边缘分割的平滑程度。实验结果证明,本文提出的模型PRA-UNet3+在对肝脏图像分割时效果明显提升,用较小的性能开销,使评价指标提升了2%以上。