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

《TransMamba：一种混合Transformer-Mamba网络用于单图像去雨》学习笔记

news 来源：原创 2025/9/16 4:03:23

paper：2409.00410

GitHub：sunshangquan/TransMamba

摘要

1、介绍

2、相关工作

2.1 单图像去雨

2.2 视觉Transformer

2.3 光谱域中的Transformer

2.4 光谱域中的图像恢复

2.5 视觉Mamba

3、方法

3.1 整体网络架构

3.2 光谱域变换块（SDTB）

3.2.1 光谱带自注意力模块（SBSA）

3.2.2 光谱增强前馈模块（SEFF）

3.3 级联双向SSM模块

3.4 重建损失与一致性损失

4、实验

4.1 实验设置

4.2 对比实验

4.3 消融实验

4.5 实际应用

4.6 限制与未来工作

摘要

现有的去雨 Transformer 采用自注意力机制，通过固定范围的窗口或沿通道维度的方式来操作，这限制了非局部感受野的利用。为了解决这个问题，我们提出了一种新颖的双分支混合Transformer-Mamba网络，命名为TransMamba，旨在有效捕捉长程的与雨滴相关的依赖关系。基于雨滴退化和背景的不同光谱域特征的先验知识，我们在第一个分支上设计了一个光谱带状Transformer模块。在光谱域的通道维度上执行自注意力操作，以增强建模长程依赖的能力。为了增强频率特定信息，我们提出了一种光谱增强的前馈模块，该模块能够在光谱域聚合特征。在第二个分支中，Mamba层配备了级联的双向状态空间模型模块，用于进一步捕捉局部和全局信息的建模。在编码器和解码器的每个阶段，我们执行双分支特征的通道级拼接，并通过通道压缩实现特征融合，从而更有效地集成Transformer和Mamba分支的多尺度信息。为了更好地重建清晰图像中固有的信号级关系，我们还提出了一种光谱一致性损失。在多个数据集和真实世界图像上的大量实验表明，我们的方法在与最先进的方法比较时具有明显的优越性。

1、介绍

雨滴去除：作为一种常见的恶劣天气，雨滴严重影响户外视觉，因此去雨是一个重要的低层视觉任务，旨在将带雨滴的图像恢复为清晰的图像。早期的工作利用物理先验来设计模型，如稀疏编码 [9]、低秩模型 [10]、高斯混合模型 [11] 等。然而，这些模型依赖于手动调整超参数，因此无法去除具有复杂外观和不同尺度的雨滴。

最近，卷积神经网络（CNN）在建模复杂模式方面超越了传统模型。为单张图像的去雨问题开发了各种CNN架构 [12]–[16]、[16]–[23]。然而，CNN无法学习长程依赖关系，因为卷积层的感受野范围有限。为了解决这个问题，Transformers [3]–[6]、[24]–[27] 被用来通过自注意力机制学习非局部表示来去除雨滴。然而，自注意力机制的二次计算复杂度限制了token选择。现有的去雨Transformer需要将自注意力限制在每个通道维度或固定大小的块中，这限制了它们充分利用长程依赖的能力。该设置还忽视了雨滴的内在一致性，这些雨滴在位置上动态变化，但在外观上密切相关。

为了解决从图像中高效去除雨滴的问题，我们提出了一种双分支混合Transformer-Mamba（TransMamba）网络。我们的模型由两个分支组成：一个Transformer分支和一个Mamba分支，每个分支在去雨过程中发挥不同但互补的作用。

雨滴由于其在光谱域中的频率特征而呈现出独特的挑战。我们观察到，光谱域中的每个token表示的是原始2D图像的正弦成分，提供了特定频带的全局分解表示。低频成分主要编码雨滴，通常表现为重复的、无纹理的区域。相反，高频成分代表具有丰富纹理的背景区域。为了利用这一观察，我们首先使用快速傅里叶变换（FFT）将图像转换到光谱域。这种变换使我们能够在光谱域中执行自注意力操作，根据频率分解更有效地捕捉全局依赖关系。通过将tokens分离到不同的频带中，我们引入了光谱带状自注意力（SBSA），它根据频带的重要性选择性地重新分配注意力。包含雨滴的低频带接收较少的注意力，而编码背景纹理的高频带则接收增强的注意力。这种选择性注意力通过减弱低频雨滴信号并保留高频图像细节，从而实现有效的雨滴去除。

为了进一步优化光谱域中的特征提取，我们引入了一个光谱增强前馈（SEFF）模块。该模块利用频域中的逐元素相乘与空间域中的卷积类似的属性。在我们的双分支架构中，我们对光谱域特征应用逐像素的权重和偏置，从而增强频率特定的信息。一个分支中的SiLU激活函数充当门控机制，控制另一个分支的输出，从而实现最佳的光谱增强。

虽然Transformer分支侧重于捕获全局信息，但我们的Mamba分支旨在增强序列一致性。该分支由双向状态空间模型（SSM）模块组成，确保图像序列之间的线性依赖关系得以保留。我们利用信号处理中的一致性度量，量化两个信号之间的线性关系，引入了一种光谱一致性损失函数。该损失函数确保重建的图像保持清晰图像中固有的线性关系，从而获得更精确和一致的去雨输出。

总之，我们的贡献有三个方面：

我们提出了一种双分支混合Transformer-Mamba网络。其第一个Transformer分支通过在光谱域tokens之间进行自注意力建模，实现全面的信息建模。第二个Mamba分支配备了双向SSM模块以增强序列一致性。
为了更好地提取频率特定的特征，我们开发了一种光谱增强前馈模块。通过引入光谱一致性损失函数，我们确保重建信号级的线性关系。
通过在各种基准测试和真实世界图像上进行大量实验，我们展示了我们的方法在去雨性能上优于最先进的方法。

2、相关工作

2.1 单图像去雨

早期的研究通过采用雨滴的各种物理先验来处理单图像去雨问题 [9]、[11]、[28]–[30]。但这些方法通常需要经验性地调整手工设计的超参数，且难以适应复杂的雨天场景。

近年来，深度学习方法（如CNN）在许多视觉任务中超越了传统方法，包括单图像去雨 [33]–[39]。许多网络设计被提出，例如编码器-解码器 [40]、细节网络 [41]、[42]、递归网络 [12]–[14]、脉冲网络 [43]、神经表示学习 [7]、生成对抗网络 [44]、多尺度结构 [16]–[18]、变分自编码器 [15]、分形带学习 [45]、渐进式学习 [16]、[22]、[23]、[46] 和半监督迁移学习 [19]–[21]。其他方法将深度学习与传统先验结合，以达到更好的去雨效果 [47]、[48]。为了使CNN能够训练，许多数据集被收集，包括合成雨滴数据集 [12]、[42]、[49] 和真实世界的雨滴数据集 [8]、[50]、[51]。尽管CNN方法在性能上超过了传统方法，但它们的感受野有限，因此很难有效地捕捉长程依赖。

2.2 视觉Transformer

最近，Transformer [52]、[53] 被引入到计算机视觉领域，并在多个任务中超过了CNN [54]–[60]。为了实现输入形状的灵活性和可扩展的推理，其他一些工作将卷积与Transformer结合 [61]–[63]。对于单图像去雨，Jiang等人 [26] 提出了将CNN与Transformer结合的去雨网络。Xiao等人 [4] 设计了一种纯Transformer模型，通过基于窗口和空间的自注意力来有效去除雨滴。为了更好地建模雨滴引起的退化关系，还提出了稀疏Transformer [5]、[6]。然而，现有的基于Transformer的去雨方法通常将自注意力限制在每个通道维度或者固定范围的窗口内，忽视了雨滴的位置是多样的，但在外观上是相似的，因此在光谱域中密切相关。为此，我们提出了一种光谱域带状Transformer，它根据背景特征在高频分量中的编码和在低频分量中嵌入的雨滴退化像素，自适应地对不同部分进行注意力分配，从而获得更优的去雨效果。

2.3 光谱域中的Transformer

一些工作探讨了结合Transformer架构和频域表示来解决磁共振超分辨率 [64]、音乐信息检索 [65]、故障诊断 [66]、视频超分辨率 [67]、[68] 和去模糊 [69] 等任务。一些超光谱图像分析工作被称为光谱Transformer，但它们与将空间RGB图像转换到光谱或频域的工作无关 [70]–[73]。Fang等人 [64] 通过计算输入图像的梯度图来嵌入高频先验，然后再将其输入Transformer。Ding等人 [66] 沿时间维度应用一维频率变换，以更好地编码长程序列。另一项工作 [67] 在自注意力模块中采用离散余弦变换，但忽略了原始图像信号的复数部分。Kong等人 [69] 将自注意力模块中空间域的张量乘法替换为频域中Query、Key和Value块之间的逐元素乘积。现有的频域Transformer要么沿时间维度计算一维频率用于处理序列数据，要么仅计算块级的频率表示，无法捕获全局空间依赖关系。因此，我们提出了图像级别的光谱带状自注意力（SBSA），用于全局特征提取，并通过光谱增强前馈模块（SEFF）更好地重新加权提取的信息。

2.4 光谱域中的图像恢复

一些工作引入了光谱域操作来解决图像恢复任务，包括超分辨率 [64]、视频超分辨率 [67]、[68]、去模糊 [69]、[74] 和一般恢复 [75] 等。然而，这些工作忽略了雨滴的光谱域先验，因此未能有效适应单图像去雨任务。

2.5 视觉Mamba

最近，状态空间模型 [76]–[79]，如Mamba [80]、[81]，展现了在长程建模中的高效性。尽管一些工作已经探索了它们在视觉任务中的能力 [82]–[85]，但其他工作指出，基于注意力的模型和CNN仍然在局部和短序列任务中优于Mamba [86]。为了进一步增强全局建模能力，现有的基于Mamba的去雨工作 [82]、[84]、[85] 都是在频域中随机地转换图像并执行计算。在本工作中，我们尝试将CNN、Transformer和Mamba结合在频域中，以使模型具备同时进行局部和全局建模的能力。

3、方法

3.1 整体网络架构

我们提出的TransMamba的网络架构，如图3所示，采用了对称的多层双分支编码器-解码器结构。假设输入的雨天图像为 $R \in \mathbb{R}^{3 \times H \times W}$ ，其中 H 和 W 分别表示图像的高度和宽度。我们使用一个 3×3 的卷积来生成图像块的重叠嵌入。在编码器-解码器的第一个分支中，我们结合多个 SDTB（光谱域变换块）来捕获复杂的特征并动态地识别雨水相关的因素。对于第二个分支，我们采用多个 CBSM（通道基础模块）来补充第一个分支 Transformer 块的缺失信息。两个分支在每一层的输出特征通过拼接和逐点卷积的方式进行合并。在同一层内，编码器和解码器通过跳跃连接互联，这与之前的研究一致 [3][4][87]，有助于增强训练过程的稳定性。每一层之间，我们使用pixel-unshuffle 和 pixel-shuffle 操作来实现特征下采样和上采样。

3.2 光谱域变换块（SDTB）

在我们的 TransMamba 中，每一层都有多个 SDTB 被级联，SDTB 的流程可以通过以下方程表示：

$F_n = F_{n-1} + \text{SBSA}[\text{LayerNorm}(F_{n-1})]$

$F_n = F_n + \text{SEFF}[\text{LayerNorm}(F_n)]$

其中，LayerNorm 表示层归一化， $F_n$ 表示第 n 层的特征。它包含两个子模块：SBSA 用于提取光谱和空间域的长程与局部依赖关系，SEFF 用于增强雨水条纹的频率特定信息。下面将详细介绍这些子模块。

3.2.1 光谱带自注意力模块（SBSA）

为了更好地捕捉动态分布的雨水降解特征，我们引入了光谱带自注意力（SBSA）模块。该模块包括一个可分离卷积序列，由点卷积和深度卷积、2D FFT 操作以及光谱带自注意力机制组成。可分离卷积用于提取 Query、Key 和 Value 的雨相关降解特征。这些特征提取后，将被送入自注意力机制以有效地加权特征令牌的值。FFT 操作将空间域的特征转换到光谱域，在光谱域中，每个像素都嵌入了具有特定频率的原始空间特征的正弦成分。由于分解特性，在光谱域中修改单个令牌将导致原图像的全局信号变化。因此，我们认为每个光谱域 token 编码了全局知识的部分依赖关系，提取这些 token 有助于捕获长程信息。通过上述特征提取过程，我们的 SBSA 实现了全局和局部动态特征的聚合。

可分离卷积序列（Separable Convolution Sequence）是指将传统卷积操作分解为多个独立步骤并按顺序执行的卷积结构，其核心目标是通过降低计算复杂度和参数量实现模型轻量化。

深度卷积：对输入特征图的每个通道独立进行空间卷积，即每个输入通道对应一个独立的卷积核（例如3x3），输出通道数与输入通道数相同。

点卷积：使用1x1卷积核进行跨通道的特征融合，调整输出通道数。

其过程可以表达为：

$[F_Q, F_K, F_V] = \text{Conv}_{1\times1}(F_n)$

$F_{i \in \{Q, K, V\} }= \text{C2R}(\mathbb{F}(\text{Conv}_{3\times 3}^d(F_i)))$

其中， $Conv^d_{3\times3}$ 是 3×3 的深度卷积， $Conv_{1\times1}$ 是点卷积， $\mathbb{F}$ 表示 2D FFT 操作，C2R 是通过 PyTorch 中的 view_as_real 操作进行的复数到实数的转换， $F_{i\in\left \{ Q,K,V \right \}} \in \mathbb{R}^{C \times H \times W}$ 是光谱域中的 Query、Key 和 Value 特征。

现有的视觉 Transformer [3]–[5] 通常依赖于固定范围的注意力或仅在通道维度上进行注意力计算，这是由于计算和内存效率的折衷。然而，这种固定设置限制了自注意力在长程范围内的自适应扩展，以便关联所需的特征。相比之下，我们观察到，由于雨水降解特征在光谱域中的低频部分具有重复的纹理，像素在不同频率上应分配不同的注意力。因此，我们提出了光谱带重排（SBR）方法，将光谱域的元素分类到不同频率的带段中，并在它们之间分配不同的注意力。

$F_{i \in \{Q, K, V\} }= \text{Reorganize}(\text{gather}(F_i; M))$

其中，Reorganize 表示重排操作，将光谱特征从 $\mathbb{R}^{C \times HW}$ 重排成 $\mathbb{R}^{Cb\times HW/b}$ ，gather 表示基于网格索引 M 收集特征，M 是通过为特定形状 H×W 计算网格索引来预处理的，中心索引表示高频，外围索引表示低频。

Reorganize 操作代替了传统 Transformer 的 Reshape 操作。

“gather&split”（收集和分割）模块将频谱特征根据网格索引进行操作，收集相关特征并分割成高频和低频两部分。

提取到的特征被送入自注意力机制，过程如下：

$F_n = \text{softmax}(\frac{F_Q F_K^T}{\sqrt{k}} ) F_V$

其中，k 是注意力头的数量。通过这种方式，我们可以有效地捕获跨频带和跨空间的长程依赖信息。

在获得光谱域注意力特征后，我们通过以下公式输出空间特征：

$F_n = \text{Conv}_{1\times 1} \left( \mathbb{F}^{-1} \left( R2C \left( \text{Scatter} \left( R(F_n); M \right) \right) \right) \right)$

其中， $\mathbb{F}^{-1}$ 是 2D 逆 FFT 过程，R2C 表示通过 PyTorch 中的 view_as_complex 实现的实数到复数操作，Scatter 表示基于索引 M 进行令牌检索的步骤，R 是维度重塑操作。

3.2.2 光谱增强前馈模块（SEFF）

先前的研究 [3]–[5] 通常依赖单一范围或单一尺度的卷积网络来增强局部上下文。然而，这些方法往往忽略了不同范围和尺度的动态分布雨水降解之间的相关性。在实践中，可以通过扩展卷积核大小或利用膨胀机制来提取多尺度信息。因此，我们将两个不同范围的深度卷积路径集成到前馈模块中，并结合光谱增强前馈模块（SEFF）进行多尺度信息增强。

SEFF 模块利用了卷积在空间域等同于在光谱域进行元素级乘积的性质，提出了光谱增强前馈（SEFF）模块。在 SEFF 中，成对的权重和偏置被调整并广播，以充当光谱域特征上的滤波器。这样，每个表示频率分量的 token 可以被自适应地增强或过滤，从而更好地分离雨相关的降解和背景信息。

在 SEFF 中，首先使用点卷积操作将通道维度扩展一个因子 $r$ ，然后将扩展后的张量传入两个并行分支。经过特征转换后，使用 3×3 卷积和扩张 3×3 卷积来增强多范围信息的提取。两个分支的特征被转换到光谱域，并通过一对可学习的权重和偏置增强光谱域信息。通过这种方式，我们增强了对不同频率分量的适应性。

该模块的流程可以通过以下方程表示：

$[F_1, F_2] = \text{Conv}_{1\times 1}(F_n)$

$F_1 = \text{Conv}^d_{3 \times 3}(F_1), \quad F_2 = \text{Conv}_{3\times 3}^{d,\text{dilated } }(F_2)$

$F_1 = W_1^{\uparrow} \odot \mathbb{F}(F_1) + B_1, \quad F_2 = W_2^{\uparrow} \odot \mathbb{F}(F_2) + B_2$

$F_n = \text{Conv}_{1\times1} ((\text{SiLU}(\mathbb{F}^{-1}(F_2)) \odot \mathbb{F}^{-1}(F_1)))$

其中， $\text{Conv}_{3\times 3}^{d,\text{dilated } }$ 是 3×3 的扩张深度卷积， $W_1,W_2 \in \mathbb{R}^{C \times H \times W}$ 是经过插值调整的光谱域权重， $B_1, B_2 \in \mathbb{R}^{C \times 1 \times 1}$ 是偏置，SiLU 是激活函数。

3.3 级联双向SSM模块

受到[79]的启发，我们提出通过级联双向SSM模块来进一步增强序列一致性。前向方向的表达可以通过以下公式表示：

$[F1, F2] = \text{Conv}_{1\times 1}(F_n)$

$F_1 = \text{SSM}\left( \sigma\left( \text{Conv1d}^d(\text{SA}(\text{Conv}^d_{5\times 5}(F1)))) \right) \right)$

$F3 = F1 \cdot \sigma\left( \text{CA}(\text{Conv}^d_{5\times5}(F2)) \right)$

看框架 $F_1$ 应该是是左边这条分支，用的应该是CA， $F_2$ 应该是右边这条分支，用的应该是SA，不知道是不是公式写错了。

其中，SA 和 CA 分别表示空间注意力（Spatial Attention）和通道注意力（Channel Attention）[94]。

反向方向的表达可以通过以下公式表示：

$[F4, F5] = \text{Conv}_{1\times 1}(\text{Flip}(F3))$

$F4 = \text{SSM}\left( \sigma\left( \text{Conv1d}^d(\text{SA}(\text{Conv}^d_{5 \times5}(F4)))) \right) \right)$

$F_o = F4 \cdot \sigma\left( \text{CA}(\text{Conv}^d_{5\times5}(F5)) \right)$

其中，Flip 是沿通道维度进行翻转的操作。

同上，感觉有问题。

通过通道翻转操作，卷积和SSM过程被从前向方向反转到反向方向。将前向和反向的 CBSM（级联双向SSM模块）级联起来，形成双向SSM模块。双向策略可以缓解SSM模块在长范围特征提取中的遗忘问题，更好地增强长距离信息的提取。

3.4 重建损失与一致性损失

为了监督TransMamba的训练，我们采用L1范数来计算去雨图像 $\tilde{B}$ 和干净背景图像 $B$ 之间的重建损失，表达式如下：

$L_{rec} = \| \tilde{B} - B \|_1$

此外，我们注意到，L1重建损失只在像素级别上调节去雨图像与真实图像之间的相似性，但忽略了从信号角度看，它们之间的全局一致性。图像中信号成分的固有关系在谱域中因雨迹退化的模式而被破坏。通过模拟原始信号关系，我们迫使退化图像的信号表示恢复为干净背景的原始正弦信号。因此，我们引入了图像谱域信号之间的一致性[95]，作为一种调节综合线性关系的手段，表达式如下：

$\mathbb{G}({\tilde{B}, B}) = \frac{\| \mathbb{F}(\tilde{B}) \bar{\mathbb{F}}(B) \|_1^2}{\| \mathbb{F}(\tilde{B})\bar{\mathbb{F}}(\tilde{B}) \|_1 \|\mathbb{ F}(B)\bar{\mathbb{F}}({B}) \|_1}$

其中， $\bar{\boldsymbol{\mathbb{F}}}$ 表示信号的共轭分量， $\| \mathbb{F}(F) \bar{\mathbb{F}}(F) \|_1$ 是信号 F 的谱密度，而分子 $\| \mathbb{F}(\tilde{B})\bar{\mathbb{ F}}(B) \|_1^2$ 表示去雨图像和干净图像信号之间的平方交叉谱密度。公式7的值在[0, 1]范围内变化。当两幅图像信号没有一致性时，公式7的值为0，而在完美线性一致性情况下，其值为1。因此，我们将谱一致性损失定义为：

$L_{coh} = 1 - \sqrt{\mathbb{G}({\tilde{B}, B} )}$

最终的损失函数表达式为：

$L = L_{rec} + \alpha L_{coh}$

其中， $\alpha$ 是谱一致性损失的权重。

4、实验

4.1 实验设置

a) 实现细节：

在我们的模型中，图3中SDTB（Spectral Domain Transformer Blocks）和CMSM（Channel-wise Multi-Scale Modules）的数量分别设置为 $N_{i \in \{1, 2, 3, 4\}}$ 和 $L_{i \in \{1, 2, 3, 4\}}$ ，其中分别设置为 (1, 3, 4, 4)。四个层次的注意力头的数量分别设置为 (1, 2, 4, 8)。初始通道数 $C_f$ 设置为36。对于SBSA（Spectral Band Splitting Attention），基于第IV-D节的消融实验，我们将带宽数 b 设置为2。关于SEFF（Spectral Encoding Feature Fusion），通道扩展因子 r 设置为2.667，预定义的权重大小为48×48。训练过程中，我们使用AdamW优化器[96]，并采用逐步训练[3]，初始批大小为8，初始patch大小为128。总训练迭代次数为300,000次。初始学习率设置为 $3 \times 10^{-4}$ ，用于前92,000次迭代，然后使用余弦退火方案[97]将学习率降低到 $1 \times 10^{-6}$ ，用于剩余的208,000次迭代。在数据增强方面，我们使用随机垂直和水平翻转。所有实验都在NVIDIA Tesla V100 GPU上使用PyTorch实现。

b) 基准数据集：

我们在五个常用的去雨数据集上进行实验，分别是 Rain200H [12]、Rain200L [12]、DID-Data [49]、DDN-Data [42] 和 SPA-Data [8]。其中，Rain200H 和 Rain200L 是Yang等人[12]提出的，分别表示含有重度和轻度雨条纹的基准数据集。每个数据集包含1800对合成训练图像和200张测试图像。DID-Data [42] 包含12000对合成训练图像，具有三种不同的雨密度级别，以及1200对测试图像。DDN-Data [42] 是通过从UCID [98]、BSD [99]和Google搜索的图像中提取1000张图像合成的14000对雨天图像。SPA-Data [8] 是一个大规模的真实世界雨天数据集，包含638,464个雨/干净图像patch用于训练，1000对图像用于测试。我们还使用了Internet-Data [8]，包括147张真实世界的雨天图像，用于定性对比。

c) 基准对比：

我们将我们的TransMamba与两种优化方法进行对比，分别是 DSC [29] 和 GMM [11]；八种CNN方法，包括DDN [42]、RESCAN [13]、PReNet [22]、MSPFN [16]、RCDNet [48]、MPRNet [23]、DualGCN [1] 和 SPDNet [2]；以及五种基于Transformer的方法，包括Uformer [87]、Restormer [3]、IDT [4]、DRSFormer [5]、UDR-S2Former [6] 和 NeRD-Rain [7]。

4.2 对比实验

a) 定量评估：

在我们的研究中，我们对合成数据集和真实雨天数据集进行了全面的比较分析，结果总结在表I中。正如所示，我们提出的方法在所有其他去雨方法中始终表现优异，特别是在PSNR方面。例如，TransMamba的性能比最先进的方法UDR-S2Former平均提高了0.32 dB。模型大小、复杂度、推理运行时以及在无参考数据集Internet-Data [8]上的无参考图像质量结果已在补充材料中报告。

由于Internet-Data [8]中缺少真实世界雨天图像的地面真值，我们通过计算无参考图像质量评估指标（包括NIQE [100]和BRISQUE [101]）对20张随机抽样的图像进行了定性比较。结果见表II。正如所见，我们的方法在这两个无参考指标上都取得了最佳图像质量。

b) 定性评估：

此外，我们在六个数据集上进行了视觉对比，去雨结果展示在图4到图7中。图4和图5展示了在合成数据集上的去雨结果，而图6和图7则展示了在真实世界数据集SPA-Data [8]和Internet-Data [8]上的去雨结果。这些视觉结果突出了我们方法在去除雨迹退化方面的有效性，包括在合成和真实世界的雨天图像中，既有小雨条纹也有大雨条纹。

c) 性能与复杂度

我们还在模型大小和计算复杂度方面将我们的方法与近期的去雨模型进行了对比。此外，我们还展示了Rain200L [12]数据集上的去雨结果（以PSNR为参考），以便便于对比这些方法的性能。复杂度的值展示在表III中。正如所示，我们的方法在模型大小和计算开销方面与之前的方法相当。值得注意的是，谱域权重是在复数空间中定义的，可以通过将它们修剪为实数空间来压缩其复杂度。我们将基于模型压缩技术，如修剪和量化，来处理模型效率的问题。总体而言，我们的方法在去雨性能和模型复杂度之间取得了良好的平衡。

4.3 消融实验

为了评估TransMamba中每个组件的有效性，我们在Rain200H数据集[12]上进行了消融实验。具体来说，我们研究了以下几个成分的影响，包括自注意力机制的选择、SBSA中通道数和频带数、前馈网络的选择，以及光谱一致性损失的权重。

a) SBSA

为了验证所提出的SBSA的有效性，我们与三个基线方法进行了比较：多维卷积头转置注意力（MDTA）[3]、基于频域的自注意力求解器（FSAS）[69]和top-k稀疏注意力（TKSA）[5]。此外，我们还探讨了通过排除SBR步骤得到的SBSA额外设置。定量分析结果如表IV所示。FSAS在频域内采用逐点积来替代常规的张量乘法，从而丧失了应由张量乘法捕捉的全局信息。尽管TKSA整合了跨通道的丰富局部信息，但可能导致在光谱域中长程依赖关系的提取不足。相对而言，我们的SBSA结合SBR能够在光谱和空间域中提取长程依赖关系，使PSNR比TKSA提高了0.47dB。

b) 通道数和频带数

为了评估SBSA中特征通道数（Cf）和光谱频带数（b）的影响，我们进行了五种不同配置的实验，分别为（72,1）、（36,2）、（18,4）、（9,8）和（4,18）。它们的乘积限制在72以内，受限于内存。结果如表V所示。可以观察到，虽然减少通道数会降低网络的学习能力，但配置（18,4）仍能提供与（72,1）类似的性能，表明SBSA的有效性。当（Cf, b）组合为（36,2）时，表现最好。

c) SEFF

为了评估所提出的SEFF，我们与四种替代方案进行了比较：（1）原始前馈网络（FN）[102]、（2）门控Dconv前馈网络（GDFN）[3]、（3）基于频域的判别前馈网络（DFFN）[69]、（4）混合尺度前馈网络（MSFN）[5]。定量结果如表VI所示。虽然DFFN能够提取频域特征，但在利用多范围空间知识和增强频域特定信息方面有所欠缺。尽管MSFN整合了混合尺度信息，但仍可能忽略了光谱域中特征的利用。通过结合FFT和跨不同范围的特征聚合，我们的SEFF模块进一步提升了性能，相较于MSFN，PSNR提升了0.14dB。

d) SSM

为了评估所提出的CBSM的有效性，我们与五种设置进行了比较：（1）视觉状态空间模块（VSSM）[83]、（2）两个前向方向级联的CBSM、（3）两个反向方向级联的CBSM、（4）前向和反向方向级联的CBSM、（5）前向和反向方向级联的CBSM（我们的方法）。如表VII所示，双向设置（4）和（5）表现最好。

e) 光谱一致性损失

表VIII展示了光谱一致性损失Lcor的有效性，并检查了不同损失权重值的影响。可以明显看出，加入Lcor后，性能得到了持续增强，当其损失权重α为5时，达到了最佳结果。

4.5 实际应用

为了进一步展示TransMamba在实际雨水去除中的应用及其对下游检测任务的提升能力，我们在图8中展示了两个示例。如图所示，所有的雨滴痕迹都被去除，帮助检测器准确识别了之前被遮挡的物体，如鞋子和手袋。

为了更广泛地评估图像去雨对下游视觉应用（如目标检测）的影响，我们利用流行的目标检测模型YOLOv3[103]评估了COCO350和BDD350数据集[16][104]上的去雨结果。表IX展示了我们的TransMamba在所有去雨方法中，获得了最佳的目标检测性能。图9中的两个实例展示了，TransMamba去雨后的图像在图像质量和检测准确度方面都表现最好。

4.6 限制与未来工作

由于TransMamba在去除雨水积聚方面的强大能力，可能会使恢复的背景过于平滑。我们计划利用扩散模型来缓解这一问题。此外，我们正在开发模型压缩和加速算法，以使雨滴去除模型能够实现实时使用。

摘要

1、介绍

2、相关工作

2.1 单图像去雨

2.2 视觉Transformer

2.3 光谱域中的Transformer

2.4 光谱域中的图像恢复

2.5 视觉Mamba

3、方法

3.1 整体网络架构

3.2 光谱域变换块（SDTB）

3.2.1 光谱带自注意力模块（SBSA）

3.2.2 光谱增强前馈模块（SEFF）

3.3 级联双向SSM模块

3.4 重建损失与一致性损失

4、实验

4.1 实验设置

4.2 对比实验

4.3 消融实验

4.5 实际应用

4.6 限制与未来工作

相关文章：