【论文阅读笔记】Multi-scale Transformer Network with Edge-aware Pre-training for Cross-Modality MR Image Syn

Li Y, Zhou T, He K, et al. Multi-scale Transformer Network with Edge-aware Pre-training for Cross-Modality MR Image Synthesis[J]. IEEE Transactions on Medical Imaging, 2023. 【开源】

论文概述

本文提出一种基于多尺度变换网络（MT-Net）的方法，用于跨模态磁共振成像（MR）图像合成。这种方法通过边缘感知的预训练和多尺度细化调整来提高合成图像的质量。核心创新包括：1) 一个边缘感知的掩码自编码器（Edge-MAE），用于预训练，以改善图像的边缘细节；2) 一个多尺度变换网络，用于细化调整，以更好地捕捉和融合不同尺度的特征；3) 引入了双尺度选择性融合模块和特征一致性模块，以进一步提升合成图像的质量和一致性。这种方法在提高跨模态MR图像合成质量方面显示出了显著的潜力，特别是在保持图像细节和结构的准确性方面。

• 研究背景和动机：作者指出，跨模态MR图像合成在医学成像中具有重要意义，尤其是在某些成像模态不可用的情况下。然而，现有的方法在处理这一任务时存在一些限制，如对边缘信息的忽视和训练不稳定性。
• 创新点概述：
  • Edge-MAE预训练策略：文章提出了一种名为Edge-MAE的预训练策略，用于改善图像合成的质量。这种策略通过随机遮盖图像的一部分来训练模型，重点关注图像的边缘信息。在预训练阶段，模型首先关注易于估计的区域（“easy patches”），然后逐渐转移到更难估计的区域（“hard patches”）。
  • 多尺度变换网络（MT-Net）：提出了一种新的多尺度变换网络，用于微调预训练模型，以适应跨模态合成任务。这个网络采用编码器-解码器架构，利用多尺度特征来提高合成图像的质量。
  • 双尺度选择融合模块：为了有效地融合来自双分支编码器的多尺度特征，文章提出了一个双尺度选择融合（DSF）模块。该模块包括空间选择、通道选择和特征融合三个关键部分。
  • 特征一致性模块：与基于生成对抗网络（GAN）的方法不同，本文采用预训练的Edge-MAE作为特征一致性模块，以增加训练稳定性并改善感知结果。
• 实验和评估：文章通过一系列实验验证了所提方法的有效性。这些实验表明，该方法在合成质量和边缘保留方面优于现有技术。
• 应用：作者强调了该技术在医学成像，特别是在阿尔茨海默病和精神分裂症等疾病的诊断中的潜在应用。

方法概述：

• Edge-Aware Pre-Training

  • Edge-Preserving Masked AutoEncoder：类似于MAE结构，但MAE由于只采用像素级损失，忽略了边缘轮廓情况，对于医学图像会影响性能，因此本文使用双decoder，以多任务方式训练，一个decoder输出还原的像素图像，另外一个输出边缘图，与Sobel算子提取的边缘对照。

  • Patch-Wise Loss：作为本文的创新点之一，这种损失函数的设计考虑到了图像中不同区域的难易程度，以及在预训练阶段对这些区域的不同处理方式，主要特点和工作原理是：

    • 
    • 基于难易程度的加权：在预训练阶段，输入图像被分割成多个8×8的小块（patches），并随机遮盖一部分。这些小块根据周围信息的丰富程度被分为“易处理”（easy patches）和“难处理”（hard patches）。易处理的小块周围有更多未被遮盖的区域，而难处理的小块则相反。
    • 动态权重分配：每个小块被赋予一个权重$\alpha$，这个权重反映了该小块的难易程度，论文中生成一个和原始图像相同size 的binary mask，对于被mask的像素标注为1，否则为0，然后根据划分的patch大小逐区域做均匀池化获得$\alpha$ 。在预训练的早期阶段，模型优先处理易处理的小块，随后逐渐转向处理难处理的小块。
    • 损失函数的设计：Patch-Wise Loss包括两个阶段的损失函数，$L_{stage1}$ 和 $L_{stage2}$。在第一阶段，易处理的小块被赋予更大的权重（通过系数$2-\alpha$），而在第二阶段，难处理的小块被赋予更大的权重（通过系数$1+\alpha$）。这种设计促进了模型在处理不同难度区域时的适应性和灵活性。
    • 加速收敛：通过这种方式，Patch-Wise Loss能够根据每个小块的难易程度动态调整训练过程，有助于模型更快地收敛，并提高整体的图像合成质量。

• Multi-Scale Fine-Tuning

• 

  • Dual-Scale Selective Fusion Module:DSF模块旨在适应性地融合来自两个独立编码器分支的特征。这两个分支分别处理不同尺度的特征，一个处理小尺度特征，另一个处理大尺度特征。

    • 三个关键组件：DSF模块包含三个主要部分：空间选择（spatial-wise selection）、通道选择（channel-wise selection）和特征融合。
    • 
      • 空间选择：在这一阶段，来自解码器的特征与两个尺度的编码器特征（一个经过上采样）进行拼接，然后通过逐点卷积操作生成两个空间注意力图。这些注意力图通过Sigmoid函数缩放到[0, 1]范围内，用于加权输入特征，以突出信息丰富的特征并抑制噪声。
      • 通道选择：在通道选择阶段，对加权输入特征进行全局平均池化，以获得通道上下文描述符。然后，结合这些描述符并通过softmax操作生成通道注意力图，用于进一步选择信息丰富的特征。
      • 特征融合：最后，通过加权结合两个分支的特征来生成融合特征图。这个融合特征图随后与来自解码器的特征进行拼接，作为下一个上采样阶段的输入。

  • Feature Consistency Module:目的是提高训练稳定性，并增强合成图像的感知质量。这种损失函数考虑了合成图像和真实图像在不同层次特征表示上的差异，有助于确保合成图像在内容和风格上与真实图像更加一致。

    • 工作原理：

      • 合成图像（记为 $\hat{y}$）和相应的缺失模态真实图像（记为 $y$）被同时输入到预训练的Edge-MAE编码器中。

      • 该模块从两个图像中提取多层次的特征。

      • 通过计算这些特征之间的差异，模块生成一个特征一致性损失（Feature Consistency Loss），用于量化合成图像和真实图像在感知层面的差异。

    • 优势：

      • 增强的稳定性：与基于GAN的方法相比，Feature Consistency Module提供了更稳定的训练过程。

      • 改善的感知质量：通过确保合成图像在特征层面与真实图像相似，可以提高合成图像的感知质量。

• 数据集和硬件

  brain tumor segmentation challenge 2020 (BraTS2020) dataset

​ ischemic stroke lesion segmentation challenge 2015 (ISLES2015) dataset

• 性能比较

  

• 方法局限性

  设计目的是学习从单个源模态到多个缺失模态图像的映射，导致无法合并来自给定患者的多个源模态图像的互补信息。