论文研读：ViT-V-Net—用于无监督3D医学图像配准的Vision Transformer

news2024/9/20 7:47:13

摘要

介绍

方法

VIT-V-Net体系结构

损失函数

图像相似性度量

变形场正则化

结果与讨论

摘要

在过去的十年里，卷积神经网络(ConvNets)在各种医学成像应用中占据了主导地位并取得了最先进的性能。然而，由于缺乏对图像中远程空间关系的理解，ConvNet的性能仍然受到限制。最近提出的用于图像分类的视觉转换器(VIT)使用了一种纯粹基于自我注意的模型，该模型学习远程空间关系以关注图像的相关部分。然而，由于连续的下采样，VIT强调低分辨率的特征，导致缺乏详细的定位信息，不适合图像配准。最近，几种基于VIT的图像分割方法被与ConvNets相结合，以提高对详细定位信息的恢复。受它们的启发，我们提出了VIT-V-Net，它连接了VIT和ConvNet，以提供3D医学图像配准。

介绍

近年来，由于在自然语言处理方面取得了巨大的成功，人们对开发基于自我注意的体系结构越来越感兴趣。、Dosovitski等人。(Dosovitski等人，2020)提出了视觉转换器(VIT)，这是第一个纯粹基于自我注意的网络，并在图像识别方面取得了最先进的性能。在这一进展之后，TransUnet(Chen等人，2021年)是在用于二维(2D)医学图像分割的预先训练的VIT的基础上开发的。

然而，医学成像方式通常产生体积图像(即，3D图像)，并且2D图像不能充分利用从3D体积获得的空间对应关系。因此，发展3D方法在医学图像配准中更为可取。在这项工作中，作者提出了一项研究，以研究VIT在体积医学图像配准中的应用。提出了一种采用混合ConvNet-Transformer结构的VIT-V-Net，用于自监督体图像配准。在该方法中，VIT被应用于运动图像和固定图像的高层特征，这需要网络学习图像中点之间的远距离关系。编码级和解码级之间的长跳跃连接被用来保持定位信息流。实验结果表明，简单地将VoxelMorph的网络结构替换为Vit-V-Net，就可以获得优于VoxelMorph和传统注册方法的性能。

方法

设f和m分别为固定图像和运动图像。我们假设f和m是单通道灰度图像，并且它们是仿射对齐的。我们的目标是预测一个变换函数φ，它将m(即m◦φ)翘曲到f，其中φ=Id+u，u表示位移矢量的流场，Id表示恒等式。图1概述了我们的方法。首先，深度神经网络(gθ)使用一组参数θ(即，u=gθ(f，m))为给定图像对f和m生成u。然后，通过空间变换函数执行翘曲(即m◦φ)(Jaderberg等人，2015年)。在网络训练过程中，比较m◦φ和f之间的图像相似度，并将损失反向传播到网络中。

图1 VIT-V-Net的方法概述和网络结构

VIT-V-Net体系结构

VIT-V-Net体系结构VIT在全分辨率体积图像中的应用导致了很大的计算复杂性。在这里，作者并没有将全分辨率图像直接送入VIT。

通过一系列卷积层和最大值池(图1中的蓝框)将图像(即f和m)编码成高级特征表示。在VIT(橙色框)中，高层特征被分成N个矢量化的 $P^{3}*C$ 块，其中 $N=\frac{HWL}{P^{3}}$ ， $P$ 表示块大小， $C$ 表示通道大小。
使用可训练的线性投影将这些patch映射到潜在的D维空间。
将可学习的position embedding添加到patch embedding以保留patch的位置信息。
将生成的补丁送入Transformer编码器，该编码器由12个交替的多头自我注意(MSA)和多层感知器(MLP)块组成
最后，对VIT的输出进行整形，然后使用V-Net风格的解码器进行解码。

图2 Vision Transformer模型概述

（请注意，编码器和解码器之间也使用了长跳过连接。网络的最终输出是一个密集的位移场。然后将其用于空间变压器中，以扭曲m）。

损失函数

本研究中使用的图像相似性度量是均方误差，以及由加权参数λ控制的扩散正则化，用于在位移场u中施加平滑。

用于训练所提出的网络的损失函数可以写成：

$L(f,m,\Phi )=L_{MSE}(f,m,\Phi)+\lambda L_{diffusion}(\Phi)$

其中，λ是正则化参数，f和m分别是固定图像和运动图像，φ表示变形场

图像相似性度量

以变形后的运动图像与固定图像之间的均方误差(MSE)作为损失函数。它被定义为：

$L_{MES}(f,m,\Phi )=\tfrac{1}{\Omega }\sum_{p\in \omega }^{}[f(p)-m o\Phi]^{2}$

变形场正则化

为了增强变形场的光滑性，使用了扩散正则化。它被定义为：

$L_{diffusion}(\Phi )=\sum_{p\in \omega }^{}||\bigtriangledown u(p)||^{2}$

结果与讨论

作者在脑部核磁共振图像配准任务中进行了实验法。使用了一个内部数据集，其中包括260个T1加权的脑部MRI扫描。数据集被分成182、26和52(7：1：2)卷，用于训练、验证和测试集。将每个图像体积随机匹配到另外两个体积以形成四对f和m，得到768、104和208个图像对。结构脑MRI的标准前处理步骤，包括颅骨剥离、重采样和仿射变换，使用Freesurfer进行(Fischl，2012)。然后，将得到的体积裁剪成相等大小的160×192×224。使用Freesurfer获得了包括29个解剖结构的标记图以供评估。

将所提出的方法在Dice Score与对称正规化(SYN)、NiftyReg以及VoxelMorph-1和-2进行了比较。正则化参数λ被设置为0.02，从结果可以看出，作者提出的VIT-V网络在Dice性能方面产生了0.1的显著收益(p值如表所示)。总之，基于VIT的注册体系结构取得了优于性能最好的注册方法的性能，证明了VIT-V-Net的有效性。

表1 将所提出的方法与其他方法进行了总体骰子比较