基于图卷积网络和有效自注意力的3D腹部器官图像分割方法

2024-09-22 00:00王川李杨魏波蒋明峰
软件工程 2024年9期

关键词:深度学习;图卷积神经网络;注意力机制;医学图像分割

中图分类号:TP399 文献标志码:A

0 引言(Introduction)

医学图像分割是将原始图像划分为不同的类别,然后提取感兴趣区域的过程。准确的医学图像分割,是医学图像分析和临床诊断中必不可少的前提与基础。医学图像分割可以为器官形状和大小测量提供视觉显示,这对于疾病的早期检测和诊断具有重要意义[1-2]。然而,受医学图像采集设备的影响,腹部器官和组织之间的灰度值差异接近,导致器官与周围组织之间的界限相对模糊。此外,不同患者的器官之间存在显著的个体差异,使得医学图像分割成为一项具有挑战性的任务。为了应对这一挑战,学者们提出了许多分割方法,包括基于阈值优化的分割方法[3]、基于机器学习的分割方法[4]和基于统计形状模型的分割方法[5]。然而,这些方法依赖于手工特征,并且特征表达能力有限。本文针对以上问题,通过采用有效自注意力提取空间与通道的特征,利用动态图卷积捕获腹部器官间的动态拓扑信息,同时有效突出腹部器官的特征,解决类间相似性和类内差异性的问题。该方法在Amos22数据集上实现了3D腹部器官的精准分割,并取得了较好的结果。

1 相关工作(Relation work)

近年来,深度学习(DL)方法被广泛用于解决医学图像分割问题,特别是基于卷积神经网络(CNN)的方法,在医学图像分割方面取得了重大的进展。例如,全卷积网络(FCN)[6]及其变体(例如U-Net[7]、U-Net++[8]等)在医学图像分割方面取得了较大的成功。此外,U形网络结构和跳跃连接也越来越受到研究人员的关注。尽管基于CNN的方法被广泛使用,但是由于其局部感知和共享权重偏差,所以网络仍然受到视野有限的影响,缺乏对远程特征相关性进行建模的能力。

近年来,许多研究试图通过使用Transformer编码器解决上述问题,Transformer是一种基于注意力的模型,最初用于自然语言处理任务中序列到序列的预测[9-10]。在医学图像分割中,CHEN等[11]首次应用Transformer改进医学图像的分割结果。CAO等[12]将Swin Transformer应用于编码器和解码器。在编码器中,实现了从局部到全局的自关注,在解码器中,特征图被上采样到与输入分辨率相同的尺寸,实现了像素级分割。MILLETARI等[13]提出了一种基于体积、全卷积神经网络的三维图像分割方法,模型在前列腺MRI图像上进行了端到端训练,并实现了一次预测整个体积的分割。ZHOU等[14]提出了一种3D Transformer,采用3种类型的注意力机制解决CNN缺乏远程建模能力的问题,基于局部和全局体积的自注意力集中于构建特征金字塔并提供大的感受野,跳跃注意力负责弥补编码器和解码器之间的差距。HATAMIZADEH等[15]将Transformer作为编码器用于学习输入体积的序列表示,有效地捕获了全局多尺度信息,并将不同分辨率的解码器特征跳跃连接到解码器,从而得出最终的语义分割。这是首个专门为3D图像提出的Transformer分割模型,其在多器官分割的BTCV(Multi-Atlas Labeling Beyond The Cranial Vault)挑战赛和医学分割十项全能数据集(MSD)上展现出了最先进的性能。

目前,腹部器官分割已成为医学影像领域的研究热点,然而大多数研究主要聚焦于单器官分割任务,例如对肝脏、肾脏和胰腺等单个器官的分割。由于不同患者多器官的外观个体差异很大,因此多器官分割比单个器官分割更具挑战性。图卷积网络(GCN)[16]通过将特征映射到拓扑图的节点和边上,能够敏感地捕获样本之间的动态变化[17],因此被广泛应用于深度学习任务中。此外,GCN通过可学习的参数,能够动态地调整不同区域特征的重要程度,这种机制有利于高效提取空间信息,进而提升模型在深度学习任务中的性能。因此,医学图像分割需要利用GCN捕获不同分割类别之间的动态拓扑关系。

2 方法(Methods)

2.1 网络结构

图1展示了本文方法的网络架构,该网络采用了经典的“U”形结构,主要分为编码器、解码器和图卷积3个部分。编码器端主要包含两个模块,即卷积模块(CNN Block)和有效自注意力模块(ET Block),卷积模块采用3D深度卷积提取特征,有效自注意力模块通过使用空间自注意力和通道自注意力对空间与通道维度上的信息进行编码,有效地学习丰富的空间通道特征表示,在每一阶段的后面都采用Patch Merging进行下采样。在网络底部,先通过区域池化模块(Region Pooling)将特征图投影为图表示,其中每个节点都代表不同的器官类型,随后进行图卷积运算,区域反池化模块(Region Unpooling)将图卷积计算后的图表示重新投影为语义图像表示。编码器和解码器之间通过跳跃连接传递信息,用于合并不同分辨率的输出,旨在恢复下采样操作期间丢失的空间信息,从而预测更精确的输出。与编码器类似,解码器也包含4个阶段,其中每个解码器阶段都包含一个上采样层,使用Patch Expanding将特征图的分辨率提高两倍,通道数量减少至上一层的50%。因此,将最后一个解码器的输出与卷积特征图进行融合,旨在恢复空间信息并增强特征表示,从而提升模型的分割性能。

3.2 实验结果对比

表2显示了本文方法在Amos22数据集上多器官分割结果的DSC值。为了验证该方法的有效性,本文将其与现有医学图像分割方法进行比较,包括Unet[7]、TransUnet[11]、nn-Unet[18]、nnFormer[14]、UNETR[15]。表2中加粗的数据表示最好的结果,*表示本文训练的结果,相关论文没有提供在此数据集上的训练结果。由表2的数据可知,应用本文方法获得的多器官分割结果的平均DSC值优于比较的方法。具体来说,本文方法的分割结果在其中12个器官中实现了最高的DSC,包括脾、左肾、右肾、食道、肝、胃、主动脉、下腔静脉、胰腺、十二指肠、膀胱、前列腺。与之前基于2D的方法(Unet、TransUnet)相比,本文方法在所有的分割部位都达到了最好的结果。值得注意的是,本文方法在对胆囊、右肾上腺(R)、左肾上腺(L)的分割中表现不佳,这可能是因为胆囊等器官在医学图像上表现的比较小,经过多次下采样后,细节信息消失不见。然而,本文方法在综合预测方面表现最好,这归功于本文提出的图卷积神经网络捕获了腹部器官间的动态拓扑信息及在研究中进行了有效自注意力机制长距离的建模。

图5显示了3个医学图像分割网络和本文方法之间的直观比较。可以观察到本文的方法比其他方法取得了更好的结果。在例子1中,由于分割的器官结构比较简单,所以本文方法与其他3种方法相比,结果差距并不是很大,只是在边缘的细节有略微改进,但是在例子3中,由于分割器官局部复杂且分割类别较多,所以其他3种方法在胃上的分割结果均出现了细微的错误。同样的现象在例子4中也可以看到,本文方法在肝脏的分割中取得了较好的结果。

为了进一步验证本文方法的有效性,本文对分割结果进行了3D重建,分割结果3D可视化如图6所示。在例子1中,本文方法和nnFormer都取得了比较好的结果,但TransUNet在三维空间中错误地分割了许多分散小点,这是因为TransUNet采用输入2D切片的方法处理3D医学数据集,割裂了数据集在空间上的信息。在例子2中可以明显看出,虽然nnFormer在大多数器官分割结果上与本文的方法相似,但是在脾分割结果的三维表面上存在小坑。综上所述,除了一些噪声和局部细节损失,本文的分割结果非常接近真实值。

如图7所示,在Amos22数据集下,本文方法展示了最快的收敛速度,而且随着迭代次数的增加,依然保持最好的分割精度。

3.3 消融实验结果

为了证明有效注意力机制(ET)、图卷积神经网络(GCN)和区域池化(Region Pooling)的有效性,本文通过逐步删除以上模块进行了详细的测试。在不同的模块组合设置下,获得了本文方法的6种变体,分别表示为“Baseline”“Baseline+ET”“Baseline+GCN”“Baseline+GCN+Region Pooling”“Baseline+GCN+ET”“Baseline+GCN+ET+Region Pooling”。如表3所示,有效注意力模块、图卷积神经网络和区域池化模块在Amos22数据集上应用后,均能够有效提升性能。本文的Baseline采用3DUnet架构,与Baseline相比,在编码器中应用有效自注意力(Baseline+ET)在DSC上提高了0.65百分点,在HD95上降低了1.94 mm。在网络底部中融合图卷积(Baseline+GCN)的结果将DSC提高了1.33百分点,将HD95降低了4.84 mm。通过在Baseline中结合区域池化和图卷积(Baseline+GCN+Region Pooling),性能比仅使用图卷积模块有了显著提高。通过将图卷积和有效自注意力(Baseline+GCN+ET)相结合,实验结果相较于使用单个模块有一定的提高。通过对3个模块的有效集成,本文方法(Baseline+GCN+ET+Region Pooling)获得了最好的结果,消融实验有力地证明了该方法融合各模块的有效性。

4 结论(Conclusion)

本文提出了一种新颖的3D医学图像分割网络模型,整体结构是在“U”形网络的基础上融合了图卷积神经网络和有效自注意力机制。有效自注意力机制在提供更大的感受野的同时,降低了一部分参数量,有效自注意力机制主要分为空间自注意力和通道自注意力,空间自注意力学习空间变换的区域特征,通道自注意力学习通道特征图之间的相互依赖关系。在编码器的底部采用区域池化提取高级语义特征,将高级图像表示投影为图表示。通过图卷积神经网络充分捕获腹部器官间的动态拓扑信息,同时有效突出腹部器官的特征。在上采样过程中,通过跳跃连接获取编码器的信息,以恢复图像至原始大小。实验结果表明,相较于以前基于2D切片的医学图像分割方法,本文方法在Amos22数据集上的表现有着较大的改进。

作者简介:

王川(1998-),男,硕士生。研究领域:医学图像处理,人工智能。

李杨(1986-),男,博士,副教授。研究领域:医学图像处理,深度学习。

魏波(1983-),男,博士,副教授。研究领域:优化算法理论,人工智能。

蒋明峰(1977-),男,博士,教授。研究领域:深度学习与优化方法,计算机图像处理。