当前位置：首页 > news >正文

《FMambaIR：一种基于混合状态空间模型和频域的方法用于图像恢复》学习笔记

news 来源：原创 2025/6/8 21:06:02

paper：(PDF) FMambaIR: A Hybrid State Space Model and Frequency Domain for Image Restoration

摘要

一、引言

二、相关工作

1、图像恢复

2、频率学习

3、状态空间模型（SSM）

三、框架

1、基本知识

2、整体框架

3、F-Mamba 模块（F-Mamba Block）

3.1 F-VSSBlock 架构：

3.2 F-MambaBlock 的核心优势：

4、Frequency-VisualStateSpaceBlock (F-VSSBlock)

4.1 F-VSSBlock 结构：

四、实验

1、实验细节

2、数据集与评估指标

3、水下图像增强

4、遥感图像去雾

5、低光图像增强

6、消融实验

摘要

随着深度学习的发展，图像恢复领域取得了显著的进展。目前的图像恢复方法主要依赖于卷积神经网络（CNN）和Transformer来获取多尺度特征信息。然而，这些方法在特征提取过程中很少有效地结合频域信息，从而限制了其在图像恢复中的表现。此外，几乎没有方法将Mamba与傅里叶频域相结合进行图像恢复，这限制了Mamba在频域上感知全局退化的能力。因此，本文提出了一种新的图像恢复模型FMambaIR，该模型利用频域和Mamba的互补性进行图像恢复。FMambaIR的核心是F-Mamba块，它结合了傅里叶变换和Mamba来进行全局退化感知建模。具体而言，F-Mamba采用了一个双分支互补结构，包括空间Mamba分支和傅里叶频域全局建模分支。Mamba负责建模整个图像特征的长程依赖关系，而频域分支则利用傅里叶变换从图像中提取全局退化特征。最后，我们使用前向反馈网络来整合局部信息，从而有助于改善恢复细节。我们在多个图像恢复任务上对FMambaIR进行了全面评估，包括水下图像增强、遥感图像去雾和低光图像增强。实验结果表明，FMambaIR不仅在性能上超越了当前最先进的方法，而且显著减少了计算复杂度。

一、引言

近年来，图像恢复已成为计算机视觉领域的一个关键研究方向，主要集中在重建或增强因噪声、运动模糊和环境因素等因素退化的图像。诸如遥感图像去雾 、水下图像增强和低光图像增强等任务都属于该领域。图像恢复在高级计算机视觉任务的推进中起着重要作用。

近年来，随着深度学习的发展，图像恢复领域取得了显著进展，催生了许多深度学习模型。之前的方法主要利用卷积神经网络（CNN）进行端到端的图像恢复。随着Transformer的快速发展，越来越多的人开始探索其在恢复任务中的表现。例如，Peng 等人首次将Transformer引入水下图像增强（UIE）任务，设计了一种基于通道的多尺度特征融合模块，增强了对退化图像的关注。UGIF-Net 提出了一个信息流引导的图像增强网络，通过注意力模块提取颜色信息。DehazeFormer改进了Swin Transformer并将其应用于遥感图像去雾领域。Trinity将先验信息集成到CNN和Swin Transformer中以估计雾霾参数。SFAN 将空间和频域信息结合用于遥感图像去雾，并设计了调制专家和解耦频率学习模块。DM-underwater提出了一个用于水下图像增强的扩散模型，可以有效提高迭代速度并减少迭代次数。尽管这些方法取得了一定的成果，但它们的计算复杂度较高，并且需要大量计算资源，这限制了它们的实际应用。

最近，结构化状态空间序列建模（SSM）已成为深度网络中的一种重要且有效的方法，尤其是通过选择性机制和高效的硬件设计进行增强。由于其在建模长程依赖任务中的优势，Mamba已在包括图像分类、医学分割、图像恢复、遥感分割和遥感变化检测等多个计算机视觉任务中得到应用。例如，MambaIR 就采用Mamba增强了图像恢复任务中的长程依赖建模。然而，它主要专注于空间域处理，可能无法完全捕捉图像中的全局退化。因此，Mamba在水下图像增强和遥感去雾等领域的应用仍需进一步探索，促使人们深入研究其在增强图像恢复网络长程建模能力中的潜力。此外，我们发现傅里叶变换具有良好的分离效果，退化信息主要存在于幅度分量中，而纹理细节信息则由相位表示。如图(a)所示，我们可视化了退化图像和清晰图像的快速傅里叶变换（FFT）图，并交换了这两幅图像的幅度。结果发现傅里叶变换在分离退化信息和细节信息方面效果良好。因此，我们将SSM和傅里叶变换结合，应用于图像恢复任务。FMambaIR通过傅里叶变换整合频域建模，有效处理全局退化特征（例如雾霾或噪声），同时保持较低的计算复杂度。

本文提出了一种简单且有效的图像恢复网络FMambaIR，基于状态空间和频域的图像恢复网络，具有局部捕获能力和高效的长程建模能力。FMambaIR的核心是频率-Mamba模块（F-Mamba），它有效地将频域特征与空间域表示相结合。空间分支利用Mamba分析复杂的空间细节；而频率分支则利用傅里叶变换从图像中提取全局退化特征，并通过频域注意力增强对全局信息的感知能力，这有助于处理不同类型的退化信息，克服了Mamba在处理频域信息方面的局限性。此外，我们引入了前馈网络（FFN），该网络有助于提取局部特征，有效去除噪声并增强图像细节，从而实现精确的图像恢复。我们在多个图像恢复任务上进行实验，包括水下图像增强、遥感图像去雾和低光图像增强。实验结果表明，我们的方法在这些任务中取得了积极的结果。此外，我们的方法比现有方法具有更低的计算复杂度，如图(b)所示，表明我们的方法有效地减少了计算资源消耗，具有更广泛的应用性。

因此，本文提出方法的主要贡献可以总结为以下几点：

1、我们提出了一种新型的频率状态空间模型（FMambaIR）用于图像恢复，能够捕捉长程依赖和局部上下文，且具有线性复杂度。

2、我们设计了一个频率-视觉状态空间块，全面融合频率和空间信息。空间分支利用Mamba分析复杂的空间细节，而频率分支从频域的角度提取全局特征，有助于处理不同类型的退化信息。

3、我们在多个图像恢复任务（如水下图像增强、遥感图像去雾和低光图像增强）上进行了实验。实验结果表明，我们的方法不仅降低了计算复杂度，而且取得了优异的成果。

二、相关工作

1、图像恢复

图像恢复是计算机视觉领域中的一个关键研究方向，随着深度学习的发展，该领域取得了显著的进展，并涌现出大量深度学习模型。SRCNN 利用卷积神经网络进行图像超分辨率，在性能上取得了令人瞩目的成果。随后的研究开发了更先进的网络架构，如残差模块、注意力机制、U-Net 和生成对抗网络（GANs）。生成对抗网络通过其对抗训练机制，在图像超分辨率、去雨和风格迁移等任务中表现出色，能够生成高质量、视觉上令人愉悦的图像。

Transformer架构在自然语言处理领域取得成功后，迅速引起了计算机视觉领域的广泛关注，成为许多图像恢复任务的基础框架。得益于其自注意力机制，它展现了强大的长程依赖捕捉能力和理解图像全局内容的能力。然而，随着输入图像尺寸的增大，该模型的计算复杂度显著增加，在处理大图像时会带来较大的计算负担。

最近，扩散模型作为一种新型的生成模型，在各种合成和恢复任务中展现出潜力。通过逐步学习去除噪声，扩散模型生成了细节丰富且自然的图像，尤其在极端退化条件下恢复图像时表现出色。尽管它们具有很大潜力，但由于迭代过程较慢且需要大量计算资源，这在一定程度上限制了它们在图像恢复中的广泛应用。

2、频率学习

傅里叶变换是一个常用的方法，用于检查信号的频率域特征。将信号从时域转换到频域可以提供信号的全局统计表示，这使得它能够有效地捕捉和描述信号的长期依赖性。近年来，许多研究者已将频率信息引入图像恢复领域。

（1）ARFFT：结合频率域与变换器（Transformer）来提升图像的感受野质量，增强了恢复结果。这表明，将频率域信息与深度学习模型（如Transformer）结合，可以提升模型的性能，尤其是在超分辨率任务中。

（2）SFMNet：通过双分支网络同时提取卷积和频率信息，使得网络能从傅里叶变换中提取全局结构。该方法通过交互模块充分整合频率域与空间域的信息，进一步提升了图像恢复的质量。

（3）FourLLIE 和 UFFC：这两种方法将傅里叶变换应用到低光图像恢复中，表明傅里叶变换对于图像的低光处理有显著优势。UFFC 提高了训练的效率，表明频率信息的处理不仅对图像质量有提升作用，也能优化训练过程。

（4）LCE 和 WF-Diff：这两个方法利用频率信息来增强全局信息的提取。LCE 通过频率自注意力模块改进了全局信息提取能力，而 WF-Diff 结合小波变换和频率残差扩散进一步提升了图像恢复的性能。

（4）DMOEAD 和 FreMamba：这些方法通过频率信息和不同的学习模块进行协同工作，进一步展示了频率学习在多任务和多目标问题中的潜力。FreMamba 应用于遥感图像超分辨率，表明频率选择模块在特定应用领域（如遥感图像）中可以发挥重要作用。

（5）Chen 等人的工作：通过傅里叶变换来进行遥感变化检测，展示了傅里叶域在结构差异检测中的优势。该方法通过自适应地融合局部和非局部结构差异图，进一步提升了检测性能。

（6）WavaMamba：结合了小波变换和 Mamba，有效地在全球范围内建模超高清（UHD）图像。

（7）DMFourLLIE ：提出了一个两阶段的多分支傅里叶低光图像增强网络，有效缓解了颜色失真和噪声问题。

在本研究中，我们引入了一个傅里叶空间注意力层，并将其与 Mamba 结合使用。

3、状态空间模型（SSM）

状态空间模型是一类用于建模时间序列数据中长期依赖性的数学模型。在计算机视觉领域，这些模型被广泛应用，因为它们能够有效地捕捉图像中的长距离依赖关系。传统的深度学习方法（如卷积神经网络）通常侧重于局部特征提取，而状态空间模型的优势在于其能够从全局角度建模数据的动态变化，特别是在序列数据和视频数据等任务中具有显著优势。

Mamba 是一种选择性状态空间模型，已经在 NLP 任务中超越了 Transformer，尤其是在推理效率和计算复杂度上。Mamba 的线性复杂度使得它在处理大规模数据时非常高效。由于这一点，它被广泛应用于包括 3D 点云、医学图像分割、视频处理等多种计算机视在本研究中，作者提出了将 Mamba 与频率信息结合并应用于图像恢复的方案。这种结合可能会在图像恢复任务中产生显著的效果，因为频率信息有助于捕捉图像的全局特征，而 Mamba 的状态空间建模能力能够增强模型对图像长距离依赖的理解。结合两者，能够在图像恢复中取得更好的表现，尤其是在需要全局信息的复杂任务中。觉任务中。Mamba 的优势在于能够有效地处理长序列数据或复杂数据结构，并且在推理速度和计算效率方面表现突出。

（1）MambaIR：MambaIR 在图像恢复任务中使用了 Mamba 模型，并通过引入残差状态空间块进一步增强了模型的恢复能力。该方法的核心优势在于能够利用 Mamba 模型强大的建模能力来处理复杂的图像恢复任务。

（2）UVM-Net：这是一种基于 Mamba 的 U-Net 架构，专门用于去雾任务。U-Net 是一种经典的图像分割网络结构，通过将 Mamba 引入 U-Net 中，进一步提升了去雾任务中的恢复能力和图像质量。

（3）ChangeMamba：这一网络专注于遥感变化检测任务，结合了 Mamba 网络的强大全局建模能力，有效检测了不同的变化。此类任务通常需要建模全局和时空关系，Mamba 在此类任务中的优势得以展现。

（4）RS3Mamba：该网络用于遥感图像的语义分割，通过协同完成模块整合 Mamba 分支与卷积分支的特征，从而提高分割精度。

三、框架

1、基本知识

状态空间模型（SSM）已广泛应用于计算机视觉任务，特别是 S4 和 Mamba，它们的核心是一个线性时间不变系统，该系统通过一个隐式的潜在状态 $h(t) \in \mathbb{R}^N$ 将一维函数或序列 $x(t) \in \mathbb{R}$ 映射到输出 $y(t) \in \mathbb{R}$ 。上述系统可以表示为一个线性常微分方程（ODE）：

$\frac{d}{dt} h(t) = A h(t) + B x(t)$

$y(t) = C h(t)$

其中， $A \in \mathbb{R}^{N \times N}$ 表示状态矩阵， $B \in \mathbb{R}^{N \times 1}$ 和 $C \in \mathbb{R}^{N \times 1}$ 分别表示投影参数。

接着，为了更好地适应视觉任务，以上过程被离散化。具体地，介绍了一个可学习的参数 $\Delta$ ，称为步长，连续参数 $A$ 和 $B$ 被离散化为 $A$ 和 $B$ 。通常使用零阶保持（ZOH）规则，其表达式为：

$A = \exp(\Delta A)$

$B = (\Delta A)^{-1} (\exp(\Delta A) - I) \cdot \Delta B$

离散化后，基于 SSM 的模型可以通过两种方式计算：线性递归和全局卷积，如下公式所示：

$h_t = A h_{t-1} + B x_t$

$y_t = C h_t$

$K = (C B, C A B, \cdots, C A^k B, \cdots)$

$y = x * K$

其中 $x$ 是输入， $K$ 表示一个结构化的卷积核。

以往的研究表明，引入选择机制可以使模型根据输入和上下文内容动态调整参数，从而提高其处理长序列和复杂依赖关系的能力。

2、整体框架

FMambaIR的整体框架如图所示。

给定一个退化的输入图像 $I \in \mathbb{R}^{3 \times H \times W}$ ，我们首先通过一个 $3 \times 3$ 卷积操作得到浅层嵌入特征 $I_{in} \in \mathbb{R}^{C \times H \times W}$ ，其中 $C$ 、 $H$ 和 $W$ 分别表示通道数、高度和宽度。这些嵌入特征 $I_{in}$ 被送入编码器进行特征提取。

投影（Input Projection）层：

投影层（Projection）在深度学习中通常指的是一种线性变换，用于将输入数据映射到另一个空间。输入投影层（Input Projection）的作用是将输入图像的像素值映射到一个更高维度的特征空间，以便模型能够捕获更多的语义信息并更好地表示图像内容。这种映射通常通过一个全连接层或卷积层来实现。输入投影层有助于将输入数据转换到一个共同的表示空间，以便进行后续的处理。

编码器主要由四个 F-Mamba 块和三个下采样阶段组成，最终得到四个编码后的特征：

$E_1 \in \mathbb{R}^{C \times H \times W}$

$E_2 \in \mathbb{R}^{2C \times \frac{H}{2} \times \frac{W}{2}}$

$E_3 \in \mathbb{R}^{4C \times \frac{H}{4} \times \frac{W}{4}}$

$E_4 \in \mathbb{R}^{8C \times \frac{H}{8} \times \frac{W}{8}}$

在每个下采样阶段，输入特征的长度和宽度都会减小，而通道数会增加。

类似地，解码器由三个上采样阶段和 F-Mamba 块组成。在每个阶段，解码器第 $i$ 层的特征 $D_i$ 会通过跳跃连接层与编码器第 $(i-1)$ 层的特征 $E_{i-1}$ 相加。类U-net 结构。

通过解码器后，我们得到输出特征 $D_3 \in \mathbb{R}^{C \times H \times W}$ 。由于水下图像增强和低光图像增强不需要改变图像的大小，我们使用一个 $3 \times 3$ 卷积来生成输出残差特征 $F \in \mathbb{R}^{3 \times H \times W}$ 。改变了通道数，将通道数从 $C$ 改成了3。

然后，将残差特征加到退化图像上，得到恢复后的清晰图像 $I' = I + F$ ，其中 $I' \in \mathbb{R}^{3 \times H \times W}$ 。

3、F-Mamba 模块（F-Mamba Block）

在图像恢复的领域中，基于 Transformer 的网络已经取得了显著进展，通常这些网络的主要模块包括注意力机制和前馈神经网络。然而，直接将选择性状态空间模型（SSM）应用于图像恢复任务并不能取得理想的效果。为了更好地利用SSM在长距离依赖建模方面的优势，本文提出了一种新的基于SSM的图像恢复网络架构，即 Frequency-MambaBlock (F-MambaBlock)。

3.1 F-VSSBlock 架构：

F-MambaBlock 中包含一个关键的新模块 F-VSSBlock，它用于替换原有的 VSSBlock（变压器模块）。F-VSSBlock 的设计考虑了图像恢复任务中的长距离空间依赖性和特征融合，通过其结构的设计来增强网络的恢复能力。

输入特征：
给定输入特征 $F_{in} \in \mathbb{R}^{H \times W \times C}$ ，即一个 $H \times W$ 大小、 $C$ 个通道的输入特征图。首先，我们通过一个 LayerNormer (LN) 进行归一化操作。这一步的作用是通过标准化输入特征的均值和方差，使得其适应后续的特征建模。

$M_l = \text{F-VSS}( \text{LN}(F_{in}) ) + s(F_{in})$

其中 $s \in \mathbb{R}^C$ 是可学习的缩放因子， $M_l$ 是通过 F-VSSBlock 获得的中间特征。

空间长距离依赖建模：
在 F-VSSBlock 中，通过引入空间长依赖建模模块来增强特征之间的相关性，能够有效地捕捉图像中的空间信息，这对于图像恢复特别重要。例如，对于一些模糊或低光的图像，捕捉远程的上下文信息（如背景和远离物体的细节）是非常关键的。

通道注意力与前馈网络（FFN）：
接下来，对 $M_l$ 进行归一化处理，再通过前馈神经网络（FFN）进行前向反馈建模。FFN 的作用是进一步捕获局部信息，并对输入特征进行更深层次的映射。

FFN 中采用了一个 $3 \times 3$ 卷积（DwConv 3×3），用于从特征图中提取局部信息。此外，通过引入 通道注意力机制（Channel Attention）来对不同通道的信息进行加权，帮助网络聚焦于最相关的特征通道。这种机制对于图像恢复中处理不同图像区域的重要性（如亮度、对比度、细节等）尤为关键。

$F_{out} = \text{FFN}(\text{LN}(M_l)) + s'(M_l)$

其中 $s' \in \mathbb{R}^C$ 是另一个可学习的缩放因子，最终输出特征 $F_{out}$ 是经过 FFN 处理后的结果。

跳跃连接：
最后，通过跳跃连接（skip connection）机制将前一层的信息与当前层的输出进行融合。跳跃连接有助于保留原始输入的细节，同时允许模型更好地学习复杂的特征，从而在图像恢复过程中减少信息的丢失。

3.2 F-MambaBlock 的核心优势：

（1）空间长距离依赖建模： F-VSSBlock 的设计能够有效捕捉图像中的长距离空间依赖性，这对图像恢复至关重要，尤其是在处理复杂的图像（如水下图像或低光图像）时，能够通过全局信息改善局部细节。

（2）通道注意力： 通过通道注意力机制，F-MambaBlock 使得网络能够自动关注图像中特征较为重要的通道，从而提高恢复的效果。不同的通道表示图像中的不同类型的特征（如纹理、边缘等），通过加权处理，能够帮助网络更好地聚焦于关键信息。

（3）跳跃连接： 跳跃连接有助于保持输入图像的低级信息，并将这些信息与网络中提取到的高级特征进行有效结合。这在图像恢复任务中尤其重要，因为它帮助网络在恢复过程中避免过多的特征损失。

4、Frequency-VisualStateSpaceBlock (F-VSSBlock)

F-VSSBlock 是 FMambaIR 的核心模块，它通过将频域信息提取与空间特征建模结合来增强图像恢复的效果。F-VSSBlock 的设计灵感来源于傅里叶变换（FFT）和选择性状态空间模型（SSM），并且该模块采用了双支路架构，分别处理输入图像的频域特征和空间特征。

4.1 F-VSSBlock 结构：

F-VSSBlock 是一个双分支模块，即 频域支路 (Frequency-branch) 和 选择性状态空间支路 (VSS-branch) 两个子模块，处理输入特征图 $X \in \mathbb{R}^{H \times W \times C}$ ，并最终将它们结合。具体过程如下：

（1）输入分割：
给定输入特征 $X$ ，首先通过 Split 操作 将输入特征图在通道维度上拆分成两部分： $X_1$ 和 $X_2$ ，每个部分的通道数为 $C/2$ 。

分别将 $X_1$ 和 $X_2$ 输入到频率域模块和 VSSBlock 中，得到相应的输出特征

${F_1, F_2} \in \mathbb{R}^{H \times W \times \frac{C}{2}}$

接着，沿通道维度连接，生成输出特征 $X' \in \mathbb{R}^{H \times W \times C}$ ，并使用 1x1 卷积进行增强。因此，F-VSSBlock 可以通过以下数学公式表示：

$X_1, X_2 = \text{Split}(X)$

$X' = \text{Conv1x1}(\text{Concat}[F(X_1), SSM(X_2)])$

（2）频域支路处理：
在频域支路中，首先将 $X_1$ 输入到 傅里叶变换（FFT） 模块中，通过傅里叶变换将图像从空间域转换到频域，获得频域特征。

然后，使用 1×1 卷积 对频域特征进行线性变换，并通过 ReLU 激活函数 和 Batch Normalization（BN） 进行进一步处理。

$F = \text{BN}(\text{ReLU}(\text{Conv1×1}(\text{FFT}(X_1)))))$

该处理过程的作用是提取图像的频域特征，并增强其频率信息，使得网络能够捕捉到图像的细节和轮廓。

然后，对频域特征的 实部和虚部 分别应用 最大池化 和 平均池化，以减少空间维度，并提取重要的频域信息。池化操作有助于减少计算复杂度，同时保留频域的关键特征。

在池化操作之后，使用 7×7 卷积 来进一步提取特征，并通过 Sigmoid 激活函数 生成空间注意力图（spatial attention map）。最后，这个空间注意力图与增强后的频域特征进行逐元素相乘（元素级加权），并通过 逆傅里叶变换（IFFT） 将其转换回空间域，得到增强后的特征图 $F_1$ 。

$F_1 = \text{IFFT}(F \odot \sigma(\text{Conv7×7}(\text{Pool}(F))))$

其中 $\odot$ 表示逐元素乘法， $\sigma$ 是 Sigmoid 激活函数。

（3）选择性状态空间支路处理（VSS-branch）：
在 VSS-branch 中，输入特征 $X_2$ 被分为两个支路进行处理：

最后，两个支路的输出通过逐元素相乘进行融合，然后通过 线性层 生成最终输出 $X_{\text{out}}$ ，并将其与原始输入 $X_2$ 相加，以得到最终的 VSSBlock 输出。

第一支路： 使用 线性层 处理输入 $X_2$ ，然后通过 SiLU 激活函数 进行非线性变换。

第二支路： 先通过 线性层，然后通过 深度卷积（DwConv） 和 SiLU 激活函数，最后输入到 SS2D 模块中进行长距离建模，并应用 Layer Normalization。

VSSBlock 的处理过程可表示为：

$X_{2_{1}} = \text{SiLU}(\text{Linear}(X_2))$

$X_{2_{2}} = \text{LN}(\text{SS2D}(\text{SiLU}(\text{DwConv}(\text{Linear}(X_2)))))$

$X_{\text{out}} = \text{Linear}(X_{2_{1}} \odot X_{2_{2}}) + X_2$

四、实验

1、实验细节

我们的FMambaIR使用PyTorch 1.13框架实现，并在配备NVIDIA RTX 8000 GPU的Linux工作站上训练。优化器采用Adam，我们设置β1=0.5和β2=0.999。训练图像随机裁剪为256×256，批量大小设置为8。网络训练1000个epoch。初始学习率设置为2×e−4，并使用余弦退火策略动态调整。训练中使用L1损失函数。

2、数据集与评估指标

数据集。我们的模型在UIEBD、和LSUI、数据集上进行了训练。UIEBD数据集包含890张真实图像，其中800张用于训练，其余用于测试。LSUI是一个真实的水下图像数据集，包含4212张图片，我们随机选择3879张用于训练，400张用于测试。此外，我们还在实际应用数据集上进行了测试，以验证FMambaIR的有效性，包括U45、、Seathru、SQUID、和Color7。

评估指标。我们引入了参考与非参考评估指标来评估我们模型的效果。对于参考指标，我们选择了峰值信噪比（PSNR）、结构相似性（SSIM）和均方误差（MSE）。对于非参考评估指标，我们选择了水下图像质量度量（UIQM）、水下彩色图像质量评估（UCIQE）、URanker和自然图像质量评估（NIQE）来评估我们的模型。