LK-CAUNet:基于交叉注意的大内核多尺度可变形医学图像配准网络
1.
2.
3.
LK-CAUNet: Large kernel multi-scale deformable medical image registration network based on cross-attention
1.
2.
3.
通讯作者:
收稿日期: 2023-06-12 修回日期: 2023-07-07 接受日期: 2023-07-14
基金资助: |
|
Received: 2023-06-12 Revised: 2023-07-07 Accepted: 2023-07-14
作者简介 About authors
程天琪(1997—),ORCID:https://orcid.org/0009-0006-8274-2287,女,硕士研究生,主要从事医学图像处理研究. 。
关键词:
Keywords:
本文引用格式
程天琪, 王雷, 郭新萍, 王钰帏, 刘春香, 李彬.
CHENG Tianqi, WANG Lei, GUO Xinping, WANG Yuwei, LIU Chunxiang, LI Bin.
近年来,基于深度学习的方法在医学图像配准领域迅速发展。试图通过学习运动图像和固定图像中变化的位移矢量场的通用表示,将运动图像映射到固定图像[10]。最典型的为VoxelMorph[11],采用UNet风格的架构估计密集的空间变形场。因只需在推理过程中进行一次前向传播,其速度比传统方法快几个数量级。此后,大量针对不同配准任务的基于深度学习的方法被相继提出[12-18]。然而,具有挑战性的应用和计算机视觉的最新进展表明,卷积操作的有效接受场受卷积核大小的限制,卷积神经网络在建模图像的远程空间关系上具有局限性。因此,自注意力机制具有大尺寸有效接受野的Transformer成为研究热点。CHEN等[19]在V-Net风格的卷积网络中采用vision transformer块提出了ViT-V-Net,并对其进行了改进,用swin-transformers取代VoxelMorph中的基本卷积块,进一步提出了TransMorph[17],使模型能够捕获远程空间信息,提高配准性能。虽然这些方法证实了transformer在图像配准中的有效性,但需大量的参数和GPU内存,大大增加了计算复杂度。
卷积神经网络在理解特征之间的空间对应关系方面能力有限,虽然基于transformer的配准模型可在一定程度上解决这一问题,但是巨大的参数量使模型在很多配准任务中难以实现。受transformer的启发,嵌入自注意力机制的模型在图像配准任务中表现出色。为提高卷积神经网络在医学图像中的配准精度,张纠等[20]提出一种双通道级联注意力网络,在配准过程中引入注意力机制进行特征权重分配,增强网络的特征表达能力。秦庭威等[21]将残差注意力机制引入动态图卷积网络,使得点云空间信息被有效利用,并减少了信息损失。虽然这些深度学习网络在单幅图像特征表示上能力较强,但在可变形医学图像配准中对图像对的特征提取和匹配仍存在局限性。如图1所示,可变形医学图像配准方法先将移动图像与固定图像融合,模拟单幅图像输入条件,再将融合后的图像输入运动-固定特征的单幅图像网络,出现混合区域内的特征失真[22-23]。这些方法因将特征提取和特征匹配过程相混合,导致混合区域特征失真和弱对齐,从而无法实现图像对之间的一一对应。最终导致关键结构缺失和配准细节较差。
图1
为了捕获运动-固定图像对之间的对应关系,提出基于交叉注意的大内核多尺度可变形医学图像配准网络(large kernel multi-scale deformable medical image registration network based on cross-attention,LK-CAUNet),如图2所示。该网络包括2个并行的特征提取子网络,各自以交叉注意的形式进行特征融合和匹配,通过基于交叉注意的融合模块,逐步融合匹配不同图像的特征表示,从而实现运动图像与固定图像对应的有效特征表示,将获得的多层次语义信息进行精细配准。同时,为匹配transformer模型的大尺寸有效接受野,在模型中嵌入可增加有效接受野的不对称并行卷积块,使其能够学习更复杂的图像结构,捕捉图像的细微变化,并阐明编码的全局相关性。此外,利用微分同胚映射计算平滑可逆的变形场。
图2
图2
基于交叉注意的特征融合与特征匹配
Fig.2
Feature fusion and feature matching based on cross attention
1 LK-CAUNet模型
1.1 整体框架
给定空间
其中,
LK-CAUNet模型的整体框架如图3所示。首先,将运动图像和固定图像分别输入各自带有大内核多尺度特征提取卷积块的特征提取子网络,以交叉注意的形式对特征进行融合和匹配,对融合的特征进行上采样后,获得运动图像与固定图像的全分辨率位移场。然后,用平方和缩放层保证最终变形的异胚性。最后,通过最小化损失函数估计变形量,将运动图像变形为与固定图像相似。
图3
1.2 LK-CAUNet模型
通过将大内核多尺度特征提取卷积块和交叉注意力模块集成至普通UNet模型,提出一种用于医学图像配准的模型。在图像配准任务中,将运动图像和固定图像分别输入各自的编码器模块,用带有大内核多尺度特征提取卷积块的编码器提取输入图像的特征,用交叉注意力模块逐步融合匹配两部分的特征表示,将融合后的特征输入解码器进行像素分类,以重构分割输出。
由于医学图像的大变形区域形变复杂,固定尺寸的卷积核难以对其进行有效的信息提取。为了避免信息之间完整性和连续性的丢失,编码器部分包含一种带有平行臂的大内核多尺度特征提取卷积块,将不同尺度卷积核提取的特征相融合,可有效弥补离散卷积核造成的信息缺失,如图4所示。在大内核多尺度特征提取卷积块中,输入数据经初始卷积后被分成不同内核大小的并行卷积分支。首先,一个核大小为3×3×3的卷积后跟随一组核大小分别为1×1×1,3×3×3和k×k×k的并行卷积。为减少并行卷积产生的(大量)特征参数,在卷积操作中应用核分解,即将N×N×N的卷积分解为3个连续的1×1×N,1×N×1和N×1×1卷积。然后,将不同支路的输出相连接。为了保留部分完整信息,在并行卷积旁增加残差连接。最后,使用ReLU激活函数修正大内核多尺度特征提取卷积块的输出。
图4
图4
大内核多尺度特征提取卷积块
Fig.4
Large kernel multi-scale feature extraction convolution block
其中,
图5
同时,针对图像配准任务,利用微分同胚映射计算平滑且可逆的变形场,从而保证运动图像变形但不丢失显著特征,即在模型输出的稳定变形场的基础上,采用7个缩放和平方层计算积分,以诱导异胚化,最终变形场可表示为
1.3 损失函数
在获取输入图像的全分辨率微分变形场后,对运动图像进行变形,以产生配准图像,用于计算损失。
无监督损失函数由均方误差(mean square error,MSE)数据项和表示变形场梯度的
其中,N表示输入的运动图像或固定图像数,
其中,
2 实验结果和讨论
2.1 实验细节
为了评估提出的配准方法,使用OASIS数据集[25]和IXI数据集(
IXI数据集的固定图像为576幅T1加权脑MRI图像,运动图像为脑部MRI数据[16],使用FreeSurfer预处理。将所有图像尺寸裁剪为160×192×224。将数据集划分为403,58和115幅,分别用于训练、验证和测试。同时,使用包含30个解剖结构的标签图评估配准性能。
在LK-CAUNet模型中,除了在大内核多尺度特征提取卷积块中设置不同大小的卷积核尺寸外,其他所有的卷积核尺寸均设置为3×3×3。用PyTorch作为机器学习库,Adam作为优化器,批处理大小设置为1,迭代次数设置为300,训练过程学习率固定为
对配准图像与固定图像进行定量和定性比较,评估不同方法的配准性能。Dice得分是一种基于重叠度的比较方法,以其作为定量评估指标。首先,使用ANTs Atropos分割方法将固定图像和配准图像分割为脑脊液、白质、灰质和背景4个部分[26]。然后,计算每部分的Dice得分,取其平均值。固定图像
2.2 结果和讨论
通过消融实验,比较UNet模型、加入大内核多尺度特征提取卷积块网络(LK-UNet)模型和本文提出的既包含大内核多尺度特征提取卷积块又包含交叉注意力模块网络(LK-CAUNet)模型在OASIS数据集上的配准性能,结果如表1所示。由表1可知,LK-UNet模型优于UNet模型,LK-CAUNet模型的Dice得分较UNet模型约提高了0.01,可见,大内核多尺度特征提取卷积块和交叉注意力模块是有效的。消融实验的可视化结果如图6所示,其中,(a)和(b)分别表示固定图像和移动图像,(c)(e)和(g)分别为3种模型的配准图像,配准结果借助紫色圆圈和橙色方框突出显示,(d)(f)和(h)分别为(c)(e)和(g)对应的变形场颜色图,颜色越深表示颜色的细腻程度越高,数值越大表示色彩过渡越平滑自然。观察发现,相较于图6(d)和(f),图6(h)大部分区域的颜色值在0.5以上,且颜色值在0.4以下的区域面积较小,因此,图6(h)对应的变形场比图6(d)和(f)对应的变形场平滑。
图6
将LK-CAUNet模型与18种具有较好配准性能的Affine、SyN[27]、NiftyReg[28]、LDDMM[29]、deedsBCV[30]、VoxelMorph[11, 22]、VoxelMorph-diff[31]、CycleMorph[16]、MIDIR[32]、ViT-V-Net[19]、CoTr[3]、PVT[33]、nnFormer[34]、TransMorph[17]、TransMorph-Bayes[17]、TransMorph-bspl[17]、UNet和TransMorph-diff[17]模型进行了比较,结果如表2所示。其中,前5种为非深度学习配准方法,为平衡配准精度与运行时间,通常根据经验设定其超参数。后13种为基于深度学习的配准方法,为保证结果的准确性,通常选择由MSE和扩散正则化项组成的损失函数进行MRI配准。
表2 不同配准方法的结果比较
Table 2
模型 | 平均Dice得分 | | J |≤0的百分比/% | 参数量/M |
---|---|---|---|
Affine | 0.386±0.195 | - | - |
SyN | 0.639±0.151 | <0.000 1 | - |
NiftyReg | 0.640±0.166 | <0.000 1 | - |
LDDMM | 0.675±0.135 | <0.000 1 | - |
deedsBCV | 0.733±0.126 | 0.147±0.050 | - |
VoxelMorph | 0.723±0.130 | 1.590±0.339 | 1.10 |
VoxelMorph-diff | 0.577±0.165 | <0.000 1 | 1.23 |
CycleMorph | 0.730±0124 | 1.719±0.382 | 0.36 |
MIDIR | 0.736±0.129 | <0.000 1 | 0.27 |
ViT-V-Net | 0.728±0.124 | 1.609±0.319 | 9.82 |
CoTr | 0.721±0.128 | 1.858±0.314 | 38.72 |
PVT | 0.729±0.135 | 1.292±0.342 | 58.80 |
nnFormer | 0.740±0.134 | 1.595±0.358 | 34.40 |
TransMorph | 0.746±0.128 | 1.579±0.328 | 46.80 |
TransMorph-Bayes | 0.746±0.123 | 1.560±0.333 | 21.20 |
TransMorph-bspl | 0.752±0.128 | <0.000 1 | 46.80 |
TransMorph-diff | 0.599±0.156 | <0.000 1 | 46.60 |
UNet | 0.727±0.126 | 1.524±0.353 | 0.28 |
LK-CAUNet | 0.828±0.138 | <0.000 1 | 9.06 |
由表2可知,LK-CAUNet模型的平均Dice得分最高,为0.828,且在变形场中几乎不产生折叠。虽然LK-CAUNet模型的参数量是CycleMorph和UNet模型的数十倍,但是其平均Dice得分提高了10%左右,且其变形场要平滑得多。同时,LK-CAUNet的参数量仅为TransMorph的1/5,但其Dice得分提高了8%左右,表明将大内核多尺度特征提取卷积块和交叉注意力模块嵌入UNet架构能产生较好的配准效果。
图7
3 结 论
针对当前配准方法存在运行时间长、无法有效保持拓扑结构、空间特征易丢失等缺点,提出了一种新型的医学图像配准网络LK-CAUNet,可以有效表示固定图像与运动图像之间的空间对应关系。通过引入交叉注意力模块,精确捕获了最终的交叉图像对应关系,实现了精细配准。同时,应用大内核多尺度特征提取卷积块,增加了模型的有效接受野,能捕获复杂图像的结构特征。对比实验结果表明,LK-CAUNet模型在可变形医学图像配准中表现优异,具有广阔的应用前景。
未来,将继续优化配准方法,进一步提高其准确性和鲁棒性,并将其应用于其他配准任务。
http://dx.doi.org/10.3785/j.issn.1008-9497.2023.06.009
参考文献
SAGE: Slam with appearance and geometry prior for endoscopy
[C]//
Self-supervised natural image reconstruction and large-scale semantic classification from brain activity
[J]. ,
CoTr: Efficiently bridging CNN and transformer for 3D medical image segmentation
[C]//
Deformable medical image registration: A survey
[J]. ,
Nonrigid registration using free-form deformations: Application to breast MR images
[J]. ,
Diffeomorphic demons: Efficient non-parametric image registration
[J]. ,
A reproducible evaluation of ANTs similarity metric performance in brain image registration
[J]. ,
Fast diffeomorphic image registration via fourier-approximated lie algebras
[J]. ,
Nesterov accelerated ADMM for fast diffeomorphic image registration
[C]//
Learn2Reg: Comprehensive multi-task medical image registration challenge, dataset and evaluation in the era of deep learning
[J]. ,
VoxelMorph: A learning framework for deformable medical image registration
[J]. ,
Topology-preserving shape reconstruction and registration via neural diffeomorphic flow
[C]//
Recursive cascaded networks for unsupervised medical image registration
[C]//
Fast symmetric diffeomorphic image registration with convolutional neural networks
[C]//
Learning a model-driven variational network for deformable image registration
[J]. ,
CycleMorph: Cycle consistent unsupervised deformable image registration
[J]. ,
TransMorph: Transformer for unsupervised medical image registration
[J]. ,
UNet vs transformer: Is UNet outdated in medical image registration?
[C]// LIAN C F, CAO X H, REKIK I, et al.
ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration
[Z]. (
基于双通道级联注意力网络的医学图像配准
[J]. ,
Medical image registration based on dual-stream cascaded attention network
[J]. ,
基于残差注意力机制的点云配准算法
[J]. ,
Point cloud registration algorithm based on residual attention mechanism
[J]. ,
An unsupervised learning model for deformable medical image registration
[C]//
End-to-end unsupervised deformable image registration with a convolutional neural network
[C]//
Non-local neural net-works
[C]//
Open access series of imaging studies (OASIS): Cross-sectional MRI data in young, middle aged, nondemented, and demented older adults
[J]. ,
An open source multivariate framework for n-tissue segmentation with evaluation on public data
[J]. ,
Symmetric diffeomorphic image registration with cross-correlation: Evaluating automated labeling of elderly and neurodegenerative brain
[J]. ,
Fast free-form deformation using graphics processing units
[J]. ,
Computing large deformation metric mappings via geodesic flows of diffeomorphisms
[J]. ,
Multi-modal multi-atlas segmentation using discrete optimisation and self-similarities
[J]. ,
Unsupervised learning of probabilistic diffeomorphic registration for images and surfaces
[J]. ,
Learning diffeomorphic and modality-invariant registration using B-splines
[C]//
Pyramid vision transformer: A versatile backbone for dense prediction without convolutions
[C]//
NNFormer: Interleaved trans-former for volumetric segmentation
[J]. ,
/
〈 | 〉 |