读论文《Multi-scale Transformer Network with Edge-aware Pre-training for Cross-Modality MRImage Synthe》

论文题目：基于边缘感知的多尺度变压器网络跨模态图像合成

论文地址：Multi-Scale Transformer Network With Edge-Aware Pre-Training for Cross-Modality MR Image Synthesis | IEEE Journals & Magazine | IEEE Xplore

项目地址：GitHub - lyhkevin/MT-Net: Multi-scale Transformer Network for Cross-Modality MR Image Synthesis (IEEE TMI)

先看论文，下一篇博客复现代码

摘要

        交叉模态磁共振(MR)图像合成可用于从给定模态生成缺失模态。现有的(监督学习)方法通常需要大量成对的多模态数据来训练有效的综合模型。然而，获得足够的成对数据进行监督训练往往具有挑战性。在现实中，我们经常有少量的成对数据，而大量的未成对数据。

        为了利用成对和非成对数据，本文提出了一种具有边缘感知预训练的多尺度变压器网络(MT-Net)，用于跨模态磁共振图像合成。具体而言，首先以自监督的方式对边缘保持掩码自动编码器(edge - mae)进行预训练，同时对每个图像中的随机掩码补丁进行图像插值，并对整个边缘映射进行估计，有效地学习上下文信息和结构信息。此外，为了提高Edge-MAE的性能，提出了一种新的基于补丁的损失方法，根据不同的掩膜补丁的impuimpu难度对其进行不同的处理。在此基础上，在后续的微调阶段，我们设计了一个双尺度选择性融合(DSF)模块，通过整合从预训练的Edge-MAE编码器中提取的多尺度特征来合成缺失模态图像。此外，该预训练编码器还用于从合成图像和相应的真值图像中提取高级特征，这些特征在训练中要求相似(一致)。实验结果表明，即使使用所有可用成对数据的70%，我们的MT-Net也可以达到与竞争方法相当的性能。

思路：编码获得每个模态的特征作为预训练，微调生成对应模拟

引言        

        磁共振成像(MRI)设置不同的扫描参数可以产生不同的模式，如T1加权(T1)、T2加权(T2)、T1加权动态对比增强(T1c)和T2流体衰减反演恢复(FLAIR)。这些对比度各有用处。

        有可能出现模态缺失，医学图像合成的目的是从给定的源模态图像中预测缺失模态(或称为目标模态)图像。

        现在的GAN 和CNN存在的一些问题，巴拉巴拉……transformer的设计植根于自关注机制，该机制有效地捕捉了输入序列元素之间的相互依存关系，以获得更好的特征表示能力

        本文提出了一种多尺度变换网络（MT-Net），通过边缘感知的自监督预训练来利用配对和非配对数据。贡献如下：

1. 提出了一种新的跨模态MRI图像合成框架，利用自监督预训练处理有限的配对数据挑战。
2. 提出了边缘保持的掩码自编码器（Edge-MAE），通过同时进行图像填充和整体边缘图估计来保留强度信息和边缘信息。
3. 设计了一种新颖的基于补丁的损失函数，通过根据各自填充的难度不同对待不同的掩码补丁，增强了Edge-MAE的性能。
4. 提出了MT-Net，通过整合预训练Edge-MAE编码器提取的多尺度特征进行微调，用于合成目标模态图像。

相关工作

• 医学图像合成：学习从给定的源模态图像到缺失模态图像的映射。现有方法大多基于条件生成对抗网络（cGANs），如Pix2Pix和CycleGAN。
• 变换器在医学成像中的应用：变换器（Transformer）最初是为自然语言处理设计的，依靠自注意力机制。在医学图像分析中，变换器已被成功应用于图像分割、分类、配准和跨模态合成。
• 自监督学习：不需要标注数据集即可学习特征表示。常见的自监督学习方法包括对比学习、自编码器方法和掩码自编码器（MAE）。

方法

该框架包括两个关键步骤:边缘感知预训练和多尺度微调。

        (a)我们使用所有可用的多模态数据进行自我监督的预训练，而不是丢弃缺少模态的样本。我们的edge - mae由一个基于变压器的编码器和两个特定任务的解码器组成，用于1)随机掩码补丁的图像输入和2)整个边缘映射估计。

        (b)对所提出的MT-Net进行了微调，以从源模态图像合成缺失模态图像。预训练的Edge-MAE编码器的前六层在微调期间保持冻结。

• 边缘感知预训练：
  • 边缘保持掩码自编码器（Edge-MAE）：由共享的基于变换器的编码器和两个特定任务的解码器组成。预训练时，输入图像被划分为一系列不重叠的补丁，随机掩码大部分补丁，剩余的未掩码补丁被投影到D维嵌入中。
  • 嵌入由标准ViT编码器编码。之后，引入可学习的掩码令牌来表示被掩码图像的每个patch。之后，两个任务特定的解码器接收全套令牌，包括未掩码补丁的潜在表示，以及可学习的掩码令牌。值得注意的是，位置嵌入被添加到令牌中，以保留补丁式的位置信息。
  • 由于原始的MAE只最小化逐像素的强度差，它不能保留结构信息，如边缘，导致输入图像中的边缘模糊。
  • 采用另一种任务特定的解码器来增强框架的边缘保持能力，该框架估计输入图像的相应边缘映射。具体来说，应用Sobel边缘检测器获得源模态图像的地面真值边缘映射。因此，我们的Edge-MAE以多任务学习方式进行预训练，同时保留了上下文和结构信息。
  • 基于补丁的损失：提出了一种新颖的基于加权ℓ1-范数的补丁级损失函数，根据其填充难度不同对待不同的掩码补丁。
  • 均方误差(MSE)平等地对待所有的补丁，而不管不同程度的补丁补齐困难。这可以归结为两个原因:a)语义信息因补丁而异(前景补丁，例如肿瘤补丁，比背景补丁包含更多的语义信息);B)随机掩蔽策略的使用也导致了不同程度的图像插入困难。

  • 

    如图2所示，将一张256 × 256 T1模态图像分割成一系列不重叠的8×8小块，其中70%的小块被随机遮挡。图2 (b)显示，绿色框中的大部分斑块是可见的，而紫色框中的大部分斑块是被遮挡的。因此，这些掩码补丁的归算困难程度各不相同。
  • 来自重遮挡区域的补丁被称为硬补丁，而来自部分遮挡区域的补丁被称为易补丁。我们的目标是在预训练的早期阶段优先考虑容易的补丁，而在随后的训练阶段分配更多的注意力给那些困难的补丁。具体来说，每个输入补丁首先被赋予一个权重α，代表输入的难度等级。为了获得每个patch的α值，基于随机掩码策略生成一个二进制掩码。
  • 例如，图2 (b)中输入图像的二值掩码大小为32 × 32，对应于patch的总数，其中值为1表示被遮挡的patch，值为0表示未遮挡的patch。然后对二值掩码进行平均池化，得到patch weight ，其中H和W分别表示输入图像的高度和宽度，P表示patch的大小。从图2 (c)中可以看出，硬斑块的α值较大，易斑块的α值较小。此外，同一patch中的所有像素共享相同的权值，所有patch的α∈[0,1]
  • 在预训练过程中，按照“易到难”的原则对掩码patch进行估算。在预训练的初始阶段，我们优先考虑从部分被遮挡的区域中插入容易的patch。在此之后，我们将注意力转移到从被严重掩盖的区域中推算这些硬斑块。提出了一种基于加权1-范数的基于patch-wise的损失算法，该算法根据不同的补全难度对掩膜patch进行不同的处理。然后，将Edge-MAE的损失函数表示为:
  • 
  • 
  • 其中x表示输入图像，y表示真实情况。λ是用来平衡不同任务的超参数。E、D分别表示编码器、用于图像输入的特定任务解码器和用于边缘图估计的解码器。S表示Sobel边缘检测器，S(y)表示输入图像的地真值边缘映射。通过对Eq.(1)中的Lstage1进行优化，其中系数2−α分配了一个权重α更大的易补片，我们对易补片进行了排序。然后，通过优化Eq.(2)中的Lstage2，我们将焦点转移到硬补丁上，其中系数1 + α分配了一个权重较大的硬补丁。提出的逐块损失算法可以加快收敛速度，简化预训练过程。
• 多尺度微调：
  • 多尺度变换网络（MT-Net）：采用编码器-解码器架构，利用预训练的Edge-MAE编码器提取的单尺度输出构建多尺度特征。

  • 

    首先利用预训练的Edge-MAE编码器的单尺度输出进行多尺度特征的构建。具体来说，我们首先对输出的小尺度特征映射进行双线性插值，从而产生更高分辨率的大尺度特征。

  • 接下来，小尺度和大尺度被馈送到我们的MT-Net的两个独立的编码器分支，其中包括多个下采样阶段。

  • 在每个阶段中，两个连续的Swin Transformer[38]层和一个补丁合并模块降低了特征地图的分辨率，同时将特征维度增加了一倍，从而实现了层次化的特征表示。与小规模分支相比，大规模分支包括一个额外的降采样阶段。解码器采用补丁扩展操作[38]和Swin Transformer层来提高特征分辨率，而跳过连接[53]允许编码器和解码器通信，从而保留空间信息

  • 双尺度选择性融合模块（DSF）：自适应地融合由双分支编码器提取的多尺度特征，并通过跳跃连接将融合的特征传递到解码器。
  • 本文提出的DSF模块由三个关键部分组成，即空间感知选择、通道感知选择和特征融合。

  • 

    我们将来自双支路编码器的多尺度特征记为和，将来自解码器的上采样特征记为。由于Fs分辨率较低，我们首先对它们进行转置卷积层，得到上采样特征。在空间选择阶段，我们分别将Fd与Fl和Fs up连接，然后将连接的特征输入到点向卷积操作中，生成两个空间注意图MA和MB。然后，利用Sigmoid函数将空间注意图缩放为[0,1]。为了在强调信息特征的同时抑制噪声特征，我们在空间注意图和输入特征之间进行空间乘法运算，从而得到过滤后的特征和。
  • 在通道选择阶段，我们对过滤后的输入特征和应用全局平均池化，以获得通道上下文描述符 ,。然后，我们将它们组合为Fp = [Pl, Ps]，然后沿着通道逻辑进行softmax操作:
  • 
  • 其中，ab分别表示相应频道注意图的第I个元素。这里，我们有a+b= 1。我们通过融合双分支特征来获得信息特征，这些特征可以用:
  • 

  • 式中，F为融合特征映射。最后，我们将F和Fd连接起来作为下一个上采样阶段的输入。

• • 特征一致性模块：使用预训练的Edge-MAE作为特征一致性模块，增加训练稳定性并提高结果的感知质量。
  • 我们采用预训练的Edge-MAE作为特征一致性模块[55]，这增加了训练稳定性并增强了感知结果。注意，我们的Edge-MAE可以很自然地作为特征一致性模块，因为它具有编码多模态和保留边缘的能力。因此，在微调过程中不再需要边缘保持损失[18]、[19]和对抗损失[24]。具体来说，合成图像y = G(E(x))和相应的缺失模态基真值y都被输入到预训练的Edge-MAE编码器中，其中x、E和G分别表示输入图像、预训练的编码器和提出的MT-Net。设Fj (y)和Fj (y)为特征一致性模块的第j变压器层的输出，该模块从y和中提取多层次特征，我们使用特征一致损失用于度量合成图像与地面真实值之间的感知差异，可以定义为

  • 

  • 式中，E为期望，l为特征一致性模块的变压器层数。在我们的框架中利用特征一致性损失使我们能够优先考虑图像之间内容和风格的相似性。此外，常规的合成图像与地面真值之间的逐像素差可以通过

  • 

  • 组合形成微调框架的总体目标函数

  • 

我们的Edge-MAE有一个具有12个变压器层的编码器，它将输入图像划分为大小为8 × 8的非重叠块。在预训练过程中，70%的图像补丁被随机屏蔽，每个图像补丁的嵌入维数为128。每个任务特定的解码器由8个转换器层组成，前3层是共享的。在解码之前，线性层将嵌入维数降低到64。本文提出的MT-Net的补丁扩展模块[38]由一个使特征维度加倍的线性层和一个扩展特征分辨率的重排操作组成。使用Sigmoid激活函数将最终预测缩放为[0,1]。

实验和结果

• 数据集：使用BraTS2020和ISLES2015数据集验证所提出的MT-Net的有效性。
• 比较方法和评估指标：与Pix2Pix、PTNet、ResViT和TransUNet等图像合成方法进行比较，采用PSNR、NMSE和SSIM等评估指标。
• 实现细节：预训练过程超过200个epoch，微调策略部分冻结预训练编码器的层。
• 合成结果：在BraTS2020数据集上进行跨模态合成任务，MT-Net在仅使用70%配对数据的情况下，合成的T2模态图像质量优于其他模型。

结论

• 贡献：MT-Net通过边缘感知预训练有效利用配对和非配对数据，即使在配对训练数据较少的情况下，也能与最先进的方法相媲美，且无需对抗训练。
• 局限性：MT-Net设计为从单一源模态学习到多个缺失模态图像的映射，未能合并给定患者的多个源模态图像的互补信息。
• 未来工作：将使框架能够接受可变数量的源模态图像进行图像合成。