从视觉基础模型到医学图像转换,利用视觉基础模型实现高效跨模态MRI合成 !

本文链接：https://blog.youkuaiyun.com/bagell/article/details/144737497

医学图像翻译对于减少临床领域冗余和昂贵的多模态成像需求至关重要。然而，目前基于卷积神经网络（CNNs）和Transformer的方法往往无法捕捉细粒度的语义特征，导致图像质量不佳。

为了解决这个挑战，作者提出了SAM-I2I，这是一个基于Segment Anything Model 2（SAM2）的新颖图像到图像翻译框架。

SAM-I2I利用一个预训练的图像编码器从源图像中提取多尺度语义特征，以及一个基于 Mask 单位注意力模块的解码器，用于合成目标模态图像。

作者在多对比MRI数据集上的实验表明，SAM-I2I优于最先进的方法，提供更高效和准确的医学图像翻译。

1 Introduction

医学影像是诊断各种疾病（尤其是在神经学和肿瘤学）的必备工具 [1]。多模态影像（如多模态MRI扫描，例如T1、T2和PD序列），可以整合不同影像序列的互补信息，从而提高诊断的准确性 [2]。不同的MRI序列可以显示不同的组织特征，有助于肿瘤、血管畸形或退行性疾病的检测 [3]。然而，获取这些不同模态需要额外的成本，包括在扫描仪中延长患者的时间。此外，由于设备限制、患者状况或长时间扫描过程中产生的运动，并非所有模态都可能可用。这些挑战推动了大量对医学影像翻译技术的需求，这些技术可以将一种影像模式合成为另一种，从而在减少多模态扫描的同时保持诊断的准确性。

早期医学图像翻译方法主要基于生成对抗网络（GANs），如 pix2pix [4] 和 CycleGAN [5]，这些方法利用对抗训练从输入源图像生成真实的目标模式图像。虽然GANs取得了有前途的结果，但它们在模式塌陷方面存在问题。此外，基于GAN的方法在从零开始训练时无法生成真实的病理区域，且数据集规模有限 [6]。自注意力机制被应用于医学图像翻译，因为这些模型利用自注意力机制捕获特征内的长期依赖关系。例如，ResViT [6] 是一种结合卷积神经网络（CNNs）和Transformer Bottleneck 的混合架构，用于跨模态MRI合成。虽然每种方法都为实现跨临床场景的最优性能提供了宝贵的洞察，但仍存在挑战。

对比学习基础模型（如对比语言图像预训练（CLIP）[7]和任意分割模型（SAM）[8]）因其通用性和广泛任务适应性而引起了广泛关注。CLIP通过在大规模图像文本对上训练，将视觉和文本表示进行对齐，使模型能够学习图像和文本嵌入之间的丰富语义关系。虽然CLIP主要用于图像分类任务，但其跨模态学习能力也使其在医疗背景[9]得到应用。SAM旨在通过学习以通用方式检测和分割目标，来处理跨任何图像域的分割任务。SAM中的图像编码器捕获了多尺度语义特征，使其成为医学图像翻译的有前景的工具，语义特征可以支持高质量的图像合成。

在本文中，作者设计了一个新的框架SAM-I2I，用于利用Segment Anything Model 2 (SAM2)实现跨模态MRI翻译。作者的方法利用SAM2中使用的预训练Hiera图像编码器作为 Backbone 模型，提供多尺度语义特征。作者设计解码器基于 Mask 单位注意力模块，有效地聚合 Backbone 模型获得的层次特征，以生成目标模态图像。在训练阶段， Backbone 模型的权重被冻结，仅训练解码器。这样的设计保留了在大数据集上训练的 Backbone 模型的表示能力。作者在一个公开的多对比MRI数据集上进行实验，并展示了与其他图像到图像转换方法相比，SAM-I2I框架的优势。

2 Methods

图1展示了SAM-I2I架构。SAM-I2I包含一个基于SAM2预训练的图像编码器和一个基于 Mask 单位注意力模块的图像解码器。预训练的图像编码器基于Hiera [10]构建，这是一种分层视觉 Transformer ，用于生成多尺度特征。对于图像解码器，使用了基于 Mask 单位注意力模块 [10]来解码跨分辨率的特征，以捕捉全局和细粒度空间细节。由于 Mask 单位注意力模块将特征分割成小局部窗口，并在每个窗口内计算注意力图，解码器还可以降低计算成本。

Hiera Image Encoder

Hiera具有四个阶段，每个阶段具有不同数量的 Transformer 块，以学习在不同尺度上的代表性特征。与其他视觉 Transformer 不同，Hiera不包含池化层，因此它逐步减小特征的空间大小，使模型能够保留更多的详细空间信息。作为SAM2中的图像编码器，它已在SA-1B [8]和SA-V [11]数据集上进行训练，这些数据集包含超过100万张图像和50,000个视频，以及1.1亿+个语义 Mask 。在这样的大型数据集上进行训练，使得Hiera图像编码器具有强大的特征表示。

Image Decoder with Mask Unit Attention

作者设计了一个图像解码器，用于根据预训练图像编码器生成的层次特征生成高质量的目标模态图像。所提出的图像解码器包含三个模块，每个模块将两个在不同分辨率下的特征集作为输入，并使用 Mask 单元注意力模块进行融合。在这里，作者决定不使用自注意力机制，因为当特征空间大小较大（例如，）时，这可能导致计算爆炸。最终卷积层将特征通道减少为1，并跟随一个 sigmoid 激活函数，将值重置在0和1之间。

在每个 Mask 单位注意力模块中，首先将特征分为几个非重叠窗口，然后在每个窗口内计算多头注意力。这样的设计可以有效降低模型的计算成本，使其在处理高分辨率特征时变得实际和高效。在模型单位注意力块中，低分辨率特征的窗口大小大于高分辨率特征，因为在低分辨率时需要全局信息以实现不同区域的交互。这有助于网络在图像中捕获更广泛的信息。相反，在高分辨率时，模型专注于局部信息以生成最终的细粒度细节。

Objective Function

在模型训练过程中，只有解码器权重被更新，编码器权重被冻结，以保持Hierarchical模型表示能力。为了训练模型，作者使用L1距离来衡量输出图像与目标图像亮度的差异，如下所示：

这里是源域图像，是目标域图像，是 SAM-I2I。仅使用可能导致过度平滑 [4]。因此，作者引入一个对抗损失定义为：

对抗生成网络采用 CLIP 图像编码器作为 Backbone 模型，通过从零开始训练可以生成质量更高的图像 [12]。最终的训练目标函数是加权组合：

其中，和是加权因子。

3 Experimental Results

作者在IXI数据集1上评估了SAM-I2I的性能。IXI数据集包含581个受试者，每个受试者都有T1加权、T2加权和PD加权的MRI图像。每个图像的空间分辨率为0.94×0.94×1.2 mm³。

作者使用80%的受试者作为训练集，剩余的作为测试集。对于每个MRI扫描，作者提取包含脑组织的轴向切片来训练网络。进行了以下四个一对一的图像转换实验：T1 → T2，T2 → T1，T1 → PD和PD → T1。

在模型训练中，作者使用Adam优化器并将学习率设置为。批量大小设置为2，所有模型均训练20个周期。和分别设置为50和1.0。合成图像使用PSNR、SSIM和NRMSE三个指标进行评估。

作者将SAM-I2I与以下基于CNN和基于 Transformer 的图像翻译模型进行了比较：UNet [13]，Pix2Pix [4]，CycleGAN [5]，以及ResViT [6]。定量结果如表1所示。作者的方法SAM-I2I在所有任务上取得了最佳性能。请注意，ResViT始终优于所有基于CNN的模型，这表明在 Bottleneck 位置应用 Transformer 块具有更强的特征学习能力对于这个数据集。

Transformer 块的长程建模特性可能有助于跨模态合成。然而，与ResViT相比，SAM-I2I在四个跨模态合成任务上的峰值信噪比（PSNR）指标平均提高了0.26 dB，这意味着相对于从头训练模型，有效地利用强大的预训练图像编码器提供的代表性特征更适合跨模态合成。此外，作者的结果表明，从其他模态学习到T1加权MRI的转换函数似乎比反向任务要简单得多。作者认为这是因为T1加权图像主要包含一般的解剖结构信息，而将T1加权图像转换为其他模态需要生成更复杂的对比度和模态特定细节。

如图2所示，不同方法生成的合成图像。作者的方法SAM-I2I的输出显示了最真实的纹理和最低的误差到真实图像。对于比较方法，合成目标模态图像往往给出了过于平滑的纹理和模糊的结构，尤其是对于ResViT。此外，对于T1 PD，CycleGAN生成的结果有一些噪声，错误图强调了大的差异。这表明CycleGAN无法有效地模拟复杂组织对比。

作者研究了不同 Backbone 模型如何影响合成的图像质量。定量结果见表2。ResNet 101 [14]和ViT-B-16 [15]的预训练权重是从CLIP在LAION-2B [16]数据集上训练得到的。ResNet 101有四个阶段，具有不同的空间分辨率，类似于Hiera模型。然而，ResNet 101中不包含任何注意力模块。

ViT-B-16首先将图像划分为几个非重叠的像素的块，并将每个图像块嵌入到一个特征向量中。这个特征向量集经过一个多头自注意力模块进行进一步的优化。

ResNet 101 始终在所有四个任务上比 ViT-B-16 获得更好的结果。这表明多尺度特征对于图像合成质量至关重要。ViT-B-16 是为了图像分类而设计的，它强调语义泛化能力，同时忽视了无用的详细局部特征。因此，ViT-B-16 可能不适合图像翻译。Hiera 模型保持了生成多尺度特征的能力，并使用自注意力机制在不同 Level 上适应性优化特征。

与 ResNet 101 和 ViT-B-16 相比，Hiera 实现了更好的性能。不同 Backbone 模型的特征图如图3 所示。Hiera 特征保留了最多的细节，而 ResNet 101 和 ViT-B-16 提取的特征分别失去了纹理信息或含有瑕疵。这证明了将预训练的 Hiera 模型作为图像翻译的图像编码器是合适的。

4 Conclusion

作者提出了SAM-I2I框架，用于跨模态MRI合成，通过利用视觉基础模型（SAM2）中编码的知识来实现。该模型使用SAM2预训练的Hiera图像编码器作为 Backbone 模型，以提取相关的层次特征。使用可训练的图像解码器，结合 Mask 自注意力单元，有效地聚合了之前提取的特征，并生成了目标模态图像。这种设计可以降低计算成本，使得网络在特征图的空间分辨率较高时仍能计算注意力图。在IXI数据集上的定量结果表明，与其他图像转换方法相比，SAM-I2I具有优势。定性结果表明，SAM-I2I可以生成具有清晰结构和真实纹理的目标模态图像。此外，使用不同的图像编码器进行消融研究，表明预训练的Hiera模型在医学图像转换方面的适用性。

目前，作者对SAM-I2I的评价仅限于跨模态MRI翻译。未来的工作将是在更广泛的医学图像翻译任务上应用SAM-I2I，如CT至MRI和PET至MRI。这些应用将使作者能够评估SAM-I2I在其他成像模式下的有效性，以满足更广泛的临床需求。此外，作者旨在研究将SAM-I2I扩展到3D体积图像合成的可行性，采用与SAM2中使用的记忆注意力机制相同的方法。