FreMIM:傅里叶变换与遮罩的图像建模在医学图像分割中的应用

代码链接：GitHub - Rubics-Xuan/FreMIM: This repo holds the official code for the paper "FreMIM: Fourier Transform Meets Masked Image Modeling for Medical Image Segmentation".

论文链接：https://arxiv.org/abs/2304.10864

收录于 WACV2024

摘要

研究界已经见证了自监督蒙面图像建模(MIM)的强大潜力，它使模型能够从未标记的数据中学习视觉表示。为了将关键的全局结构信息和局部细节信息结合到密集预测任务中，我们将视角转移到频域，提出了一种新的基于mimm的自监督预训练框架FreMIM，以更好地完成医学图像分割任务。在观察到详细的结构信息主要存在于高频成分中，而低频成分中高层次语义丰富的基础上，我们进一步在预训练阶段引入多阶段监督来指导表征学习。在三个基准数据集上进行的大量实验表明，我们的FreMIM比以前最先进的MIM方法具有优越的优势。与从头开始训练的各种基线相比，我们的FreMIM可以持续地为模型性能带来相当大的改进。

背景

Masked Autoencoders (MAE)[26]，它通过掩盖图像中的部分区域并重建它们来预训练模型。MAE仅将原始像素作为重建目标，主要依靠局部特征表示，而没有充分利用全局信息。

因此，为了在训练样本有限的情况下充分发挥基于mima的医学图像分割方法的潜力，如何在获取全局信息的同时尽可能保留详细的局部特征成为关键问题。

正如之前的许多研究[5,7,14,30,45]所示，详细的纹理信息主要存在于高频分量中，低频分量中包含丰富的全局信息。根据这一观察，一个直观的解决方案将是探索MIM与傅里叶变换相结合的强大潜力。

由于同一器官的医学图像本质上对应相似的特征，我们进行了困难的跨域重建任务，避免了走捷径的学习，实现了较强的表征能力。

同时，借鉴前人[49]的研究结果，即详细的结构信息主要存在于高频分量中，而低频分量中具有丰富的高级语义，本文提出的双边聚合解码器对原始图像依次进行傅里叶变换，并在变换后的傅里叶谱上使用低/高通滤波器，得到预期的重构目标。

相关工作

Masked Image Modeling

通过重建图像的被屏蔽部分，模型可以学习到有利于各种视觉下游任务的信息特征表示。

之前使用的随机掩蔽策略是粗糙的，可能会在无用的背景上造成计算浪费。考虑到医学图像中信息前景和无用背景的区分，我们设计了前景像素之间的掩蔽策略，以获得更有效的掩蔽，帮助模型更好地进行表征学习。此外，我们的方法可以摆脱对特定模型结构的预训练范式的依赖，并持续提高模型性能，这与以往的工作不同(例如，Swin Transformer和基于cnn的模型不能直接与MAE集成)。

Fourier Transform

[43]利用快速傅里叶变换(Fast Fourier Transform, FFT)替代原始Transformer中的自关注模块，以较低的计算成本成功获取全局信息。[29]为傅里叶频谱监督设计了一种新的焦频率损失，以提高流行的图像生成模型的性能

我们对原始图像进行随机掩码，重建图像频域中的傅里叶谱去帮助模型以跨域掩码重建的方式学习更广义的全局表示。此外，还提出了结合FFT特定特性(即高通和低通频率分量)的多级监督，以更好地指导不同阶段之间的模型表示学习。

贡献

1） 我们首次研究了在医学图像分割任务中利用频域掩膜图像建模的强大潜力。被提议的
FreMIM是一个通用的自监督预训练框架，可以与不同的模型体系结构集成
(即cnn和transformer)。

2）通过设计一个多阶段监督方案和一个设计良好的双边聚合解码器，我们提出了一种新的跨域掩码重建框架，用于掩码图像建模范式。

3）提出了一种简单而有效的前景像素间（foreground pixels）masking 策略，作为原始随机掩蔽像素策略的更好替代方案，为后续的自监督表示学习提供更精确和信息丰富的掩蔽

方法

Preliminary: Fourier Transform

由于离散傅里叶变换(DFT)在我们提出的方法中起着至关重要的作用，我们首先简要回顾了作为传统信号分析不可或缺的技术的二维DFT。给定一个2D信号F∈RW×H，其对应的2D- dft可以定义为:

其中F(h,w)表示位于F中(h,w)的信号，u和v是傅里叶频谱中水平和垂直空间频率的指标。相应的，二维逆DFT (2D- IDFT)表示为:

DFT和IDFT都可以使用它们的快速版本FFT算法进行加速[39]。对于具有多种模态的医学图像，傅里叶变换在每个通道上独立进行。此外，如先前的研究[5,7,14,30,45]所示，图像的详细结构纹理信息主要存在于傅里叶谱的高频部分，而全局信息则丰富于低频部分。图2展示了这一有趣特征的可视化。

The Proposed FreMIM

Overall Architecture

给定空间分辨率为H×W和C通道(模态数)的输入医学图像切片X∈RC×H×W，首先对原始图像采用所提出的前景掩蔽策略生成掩蔽图像。然后，通用编码器(即根据各种预训练要求，cnn和Transformers编码器都可以很容易地集成到我们的框架中)将被屏蔽的图像作为输入，通过分层结构捕获被屏蔽的视觉特征。然后，将不同阶段的编码特征表示联合馈送到我们精心设计的双边聚合解码器中，逐渐产生具有低级细节信息和高级语义表示的重构傅立叶谱。

通过顺序应用傅里叶变换在原图片上，在转换后的傅里叶频谱上采用低/高通滤波器获取期望重建目标，将重建损失应用于重建频谱与期望低/高通频谱目标的相似度，实现端到端低阶和高阶表示的有益多阶段监督方案。

Masking Strategy

医学图像的前景和背景像素分布极不平衡。因此，随机选择医学图像的空间位置，不可避免地会导致生成的掩模大部分覆盖背景像素，保留了太多物体的前景像素，不利于模型的重建能力。
为此，我们提出了一种简单而有效的前景掩蔽策略来解决这种不均匀分布问题。

具体来说，给定一个初始化为零的二进制掩码M∈{0,1}H×W，其在每个空间位置的值取决于对应的像素值是否属于前景。如果一个像素属于前景区域，在自监督预训练中，它将被滤除作为被遮罩的候选点之一。

由于医学图像通常由不同的通道组成，每个通道强调不同的前景区域，因此我们将它们的重叠部分作为最终的掩码区域。整体前景掩蔽策略可以定义为↓，式中⊙为Hadamard积，Pn(x, y)表示对应位置(x, y)的特定像素值，Mn表示生成的特定图像模态Mn的掩模。

M和XM分别表示原始图像的最终掩码，以及将被掩码后的图像馈入模型进行后续重建任务。

 

Generic Encoder

至于我们框架中编码器的选择，由于我们的逐像素前景掩蔽策略，FreMIM不限于任何特定类型的结构。网络编码器以上述掩膜图像为输入，对掩膜图像切片进行分层结构的逐步编码，产生不同层次(即从底层细节信息到高层语义)的特征表示。

Multi-stage Supervision Scheme

对有效SSL范例的期望是引导可视化主干通过分层结构学习不同级别的所需表示。根据这种直觉，我们提出设计一个多阶段监督方案，以充分监督分层阶段的表示学习。

图像的高阶和低阶信息分布在傅里叶谱的不同频段。因此我们提出分别利用低通和高通傅里叶频谱作为监督信号(即重构目标)。最直观的方法之一是利用相同的高通傅立叶谱直接监督多个low level stages ，反之亦然。然而，这种直观的方式主要有两个缺点。一方面，违背了模型在各个低级阶段学习的初衷，因为不同低级阶段学习到的特征表征自然不应该是相同的，而应该是不同的。另一方面，这种监督方法过于直接和简单，没有充分利用层次结构捕获的多阶段特征之间的相关性来帮助模型更好地执行MIM pretext 任务。

为此，我们提出了一种设计良好的双边聚合解码器 bilateral aggregation decoder，以更好地解决频域重构任务，进一步帮助编码器学习更广义、更有意义的特征表示。具体而言，在所提出的双边聚合解码器内部，将不同阶段的编码特征分别以自下而上和自上而下的方式收敛到最低阶段(即最大空间分辨率)和最高阶段(即最小空间分辨率)。换句话说，BAD将不同阶段的特征图分别聚合为最低和最高分辨率。具体来说，对于ViT，第4层、第8层和第12层的特征映射分别上采样8次、4次和2次，然后按照UNETR中的反卷积模块馈送到BAD。为了明确起见，每个相邻阶段的捕获特征将被馈送到卷积块中，以实现空间分辨率和通道维度的严格对齐，可以表示为↓其中Ahigh和allow分别表示双边聚合的高级和低级特征表示，C、Dc和Cat分别表示卷积块、反卷积块和拼接操作，Si表示阶段i输出的特征映射。

然后，通过引入的频率映射块(如图3所示)，将最低级和最高级的聚合特征表示映射到频域，然后通过低通和高通滤波器，得到所采用重构损失对应的高通和低通预测谱。具体来说，频率映射块(FMB)由2D-DFT、频域感知器(FDP)和2D-IDFT组成，其计算公式为↓ 其中DFT和IDFT分别表示快速傅里叶变换和快速傅里叶反变换。W和b都是可学习的参数，⊙是哈达玛积。通过这种方式，利用傅里叶变换的独特特性，构建了一个强大的SSL跨域重建框架。

 

虽然这样的跨域重建任务比域内重建任务更加困难，但它也可以帮助模型学习更鲁棒的特征表示

Pre-training Strategy

Frequency Loss

为了缓解不同频带频谱间的权重不平衡，便于对困难频带进行重构，我们采用焦频损失[29]作为损失函数lfrequency，对低频和高频映射进行权重梯度更新，定义为↓式中f(u, v)为空间频率坐标(u, v)的预测2D-DFT，而f^(u, v)为其对应的Ground truth。γ(f, f^)计算实际值和预测值之间的欧几里得距离的平方作为它们的频率距离。

ω是给定位置的频谱权矩阵，它抑制了简单频率的权值。计算公式如下，其中β是灵活性的缩放因子(默认β=1)。

Overall Loss

在预训练期间，我们的FreMIM通过从高通和低通频率中求解内容格式塔来学习表示:式中FH和FL分别为高通和低通频率滤波器。T表示原始图像。如图3所示，P low由最高级得到，而P high则相反。α是高级语义信息分支的权重(默认α = 3)。

实验

数据集：

BraTS 2019， [2,3,36],

ISIC 2018 [16,48] 

ACDC 2017

实验结果

总结

在本文中，我们首次研究了基于频域的MIM在医学图像分割任务的预训练深度学习模型上的强大潜力。针对二维医学图像分割问题，提出了一种利用傅里叶谱中丰富的全局信息和局部细节的新框架FreMIM。不同于以往MIM方法的传统范式，在频域实现重构使框架具有更强的表示学习能力。此外，多阶段监督方案充分利用了不同频段所包含的特定特性，大大提高了分割性能。在三个基准数据集上进行的综合实验，定量和定性地验证了我们的FreMIM的有效性，显著提高了从头训练的基线分割性能，并显示出优于最先进的自监督方法的优势。