paper:FreqMamba: Viewing Mamba from a Frequency Perspective for Image Deraining
GitHub:GitHub - aSleepyTree/FreqMamba
目录
摘要
1、介绍
2、相关工作
2.1 图像去雨
2.2 频率分析
2.3 状态空间模型
3、方法
3.1 动机
3.2 预备知识
3.2.1 傅里叶分析与数字图像
3.2.2 状态空间模型
3.3 模型架构
3.3.1 FrepSSM块
3.3.2 数据依赖性降解先验注意力图
3.4 损失函数
4、实验
4.1 数据集与实现
4.2 对比实验
4.3 消融实验
4.4 在其他任务上的扩展
摘要
由于雨条纹的影响,图像通常会丢失感知所需的重要频率信息,图像去雨的目标是解决这一问题,这需要对全局和局部退化进行建模。最近的研究显示,Mamba在利用局部相关性来感知全局和局部信息方面具有高效性和有效性。然而,目前很少有研究尝试将Mamba与频率分析结合起来应用于图像去雨,这限制了其感知基于频率建模(例如傅里叶变换)的全局退化能力。
在本文中,我们提出了一种名为FreqMamba的高效且有效的范式,它利用Mamba和频率分析之间的互补性来解决图像去雨问题。该方法的核心在于从以下两个方面扩展Mamba与频率分析的结合:通过频率带扩展来挖掘频率相关性,并结合傅里叶变换建模全局退化。
具体来说,FreqMamba引入了三种互补的交互结构:空间Mamba、频率带Mamba和傅里叶全局建模。频率带Mamba将图像分解为不同频率的子频带,以从频率维度进行二维扫描。此外,利用Mamba独特的数据相关特性,我们使用不同尺度的带雨图像为网络提供退化先验,从而促进高效训练。
大量实验结果表明,我们的方法在视觉和定量评估(例如PSNR和SSIM)上均优于当前最先进的方法。
1、介绍
背景与问题
雨天拍摄的图像因空气中的雨滴导致对象细节和对比度显著下降,不仅视觉效果不佳,还会损失重要的频率信息。这种降质对图像分类和目标检测等高级视觉任务造成了严重的不良影响。因此,图像去雨是低级视觉领域中的一项重要任务。然而,从严重雨滴退化的图像中恢复清晰图像非常困难,主要因为雨滴与背景的复杂耦合以及重要感知频率信息的丢失。
传统方法依赖先验知识来分解雨滴与背景层的物理属性,并引入各种方法区分雨条纹和干净图像。然而,这些基于特定观察的先验知识可能无法可靠地建模图像的内在特性或在物理模型中估计传输图。深度学习的出现为去雨技术带来了新的方向,许多方法从不同角度提升了去雨性能,尤其是有效的全局降质建模在解决复杂的图像去雨挑战中至关重要。例如,Transformer中的注意力机制通过建模图像内部相关性取得了巨大成功,但由于其二次复杂度,注意力机制在处理大图像时面临扩展性挑战。
最近,一种改进的结构化状态空间序列模型(S4)结合选择性扫描机制(Mamba)脱颖而出。它通过线性复杂度建模长序列关系,特别是Mamba的选择性方法能够显式建立图像块间的相关性,从而引导清晰区域到退化区域的映射。然而,值得注意的是,Mamba处理视觉任务的方式本质上基于像素序列,尽管这允许建模长距离依赖,但其一维序列的选择性建模限制了在全局降质感知(如傅里叶变换)中的能力。为此,我们尝试从频率分析的角度扩展Mamba以应用于图像去雨。
提出方法:FreqMamba
我们将两种典型的频率分析技术整合到方法中:傅里叶变换(FT)和小波变换(WPT)。众所周知,傅里叶变换具有有效的全局建模能力,并在降质分离中表现良好。这种独特视角对于图像去雨任务尤为重要,因为理解整个图像结构显著提高了恢复质量。然而,不同域之间存在固有差距,简单地将傅里叶变换与Mamba结合可能显得生硬。
为实现无缝过渡,我们需要一个中间状态。小波包变换将图像分解为不同空间域中的子频带,同时包含频率分析和空间信息。利用小波域作为空间域与傅里叶域之间的中间状态,我们建立了更平滑的过渡,从而增强分析的整体有效性。
FreqMamba架构
基于频率分析感知全局降质的潜力以及Mamba捕捉空间域局部相关性的能力,我们引入了FreqMamba。该高效且有效的范式利用了Mamba与频率分析之间的互补性。方法的核心在于三分支结构:
空间Mamba:对原始图像特征操作,提取细节和图像内相关性,提供关键的降质模式信息。
频率带Mamba:利用WPT将输入特征分解为跨越不同频率带的特征光谱。通过将这些低分辨率特征按频率顺序重构为原始尺度,进行从低到高频的频率扫描和逆向扫描。这种策略不仅丰富了模型的分析广度,还充当了空间与傅里叶域之间的桥梁,为建模提供了新视角。
傅里叶建模:借助傅里叶变换对输入进行全局分析,捕捉影响图像的整体降质模式,提供频率谱的全景视图。这种全局建模能力有助于全面理解和缓解降质效应,从而更彻底地去除雨条纹,确保清晰的视觉效果。
集成与提升上述分支共同形成了一个坚实的架构,解决图像去雨的挑战。空间Mamba分析复杂的空间细节,傅里叶建模分支提供全局视角,而频率带扫描Mamba探索频率维度,为二维建模提供新思路。此外,借助Mamba的独特数据相关特性,我们将该方法应用于不同尺度的雨天图像,生成基于降质先验的注意力图。这些注意力图随后集成到骨干网络中,以辅助高效训练。
2、相关工作
2.1 图像去雨
图像去雨经历了显著的发展,从早期的基于模型的方法转向先进的数据驱动技术。最初,基于模型的方法通过手工设计的特征和物理先验知识将雨条纹与背景分离。这些方法虽然有一定的洞察力,但通常难以应对复杂的雨滴模式和多样的现实场景,因此在实际应用和性能方面存在局限性。深度学习的出现为图像去雨开启了新时代,数据驱动的方法表现出从数据中提取和学习特征的卓越能力。卷积神经网络(CNN)的引入标志着一个重要的进展,使得在各种图像中更加细致和自适应地处理雨条纹成为可能。此外,结合注意力机制的架构的发展进一步提升了识别和去除雨成分的能力,解决了之前在模型泛化和细节保留方面的不足。在本研究中,我们提出了一种新的基线,并采用了创新的三重交互模块,旨在联合提高去雨性能。
2.2 频率分析
傅里叶变换(FT)是频域分析中的一项基础技术,它能够将信号转换到一个可以更容易分析其全局统计特性的域。这一能力被广泛应用于各类计算机视觉任务中。通过有效地建模全局域信息,傅里叶变换在多个领域的进展中发挥了重要作用。在图像恢复方面,FECNet展示了傅里叶特征幅度在隔离全局亮度成分方面的有效性,从而改善了图像的美学效果和清晰度。同样,FSDGN揭示了傅里叶特征的幅度如何作为图像去雾任务中全局雾霾信息的关键指标。尽管取得了这些进展,傅里叶变换在信号处理中的固有限制表明其潜力尚未完全开发,暗示这些方法的效能可以进一步增强。除傅里叶变换外,小波变换(WT)也是一种用于分析信号和图像的数学工具,提供了与傅里叶变换互补的视角。与傅里叶变换擅长捕获频率信息不同,小波变换提供了多分辨率分析,特别在检测和表示信号中的局部变化方面具有较高的有效性。
2.3 状态空间模型
状态空间模型(SSMs)近年来因其能够有效建模长程依赖关系,并且在序列长度变化时表现出线性扩展性,受到了广泛关注。S4的基础性工作为深度状态空间建模奠定了基础,展示了有希望的线性扩展特性。最近的创新通过整合门控单元增强了SSM的能力。此外,Mamba,一种数据依赖型的SSM,具有选择机制和高效的硬件设计,作为一个突出表现者,在自然语言处理任务中超越了Transformer,同时保持输入长度的线性扩展性。SSM的适用范围不仅限于自然语言处理(NLP),一些开创性的研究已将Mamba应用于各种计算机视觉任务,包括图像分类、生物医学图像分割等。
3、方法
3.1 动机
降质图像,特别是那些受到雨滴影响的图像,遭受着全球降质和局部细节丧失的问题。基于频率的方法被用来减轻图像整体质量的下降,利用雨滴效应与傅里叶幅度谱之间的显著相关性。然而,由于频率域与空间域之间固有的差距,频率域操作的全局性使得无法建模空间域的局部依赖关系。Mamba是一种具有选择扫描机制的状态空间模型,它通过线性复杂度巧妙地促进序列间交互的显式建模。应用于视觉任务时,它能够有效地建立二维图像区域之间的相关性。
认识到基于频率的方法和Mamba模型在处理图像降质的不同方面上的互补优势,我们引入了FreqSSMBlock。这个新颖的结构具有三分支设计,能够平滑地过渡于全球降质的修正与局部内容的细化之间。进一步利用Mamba独特的数据依赖特性,我们在不同尺度下应用降质图像来推导注意力图,并将其添加到编码器阶段,从而提高训练效率。在此基础上,我们提出了FreqMamba,一种专门设计的解决方案,旨在解决单幅图像降雨任务。
3.2 预备知识
3.2.1 傅里叶分析与数字图像
Fourier Transform(傅里叶变换)是一种将信号从原始域(通常是时间或空间)转换到频率域的数学技术,反之通过逆傅里叶变换(iFT)也可以进行转换。对于一个单通道图像 (大小为 ),其离散傅里叶变换(DFT)定义如下:
这一过程将图像的空间特征转换为复数分量,展示图像的频率成分。频率域表示 可以分解为实部和虚部,从而得到幅度谱和相位谱,描述图像的频率内容:
另一个工具是离散小波变换(DWT),它将图像分解为四个子带,分别表示低频近似和高频细节:
每个子带对应不同方向的细节,且每个子带将原图像的尺寸减少一半。与此不同的是,小波包变换(Wavelet Packet Transform)通过在每一级进一步分解所有子带,提供了更详细的频率内容分析。
3.2.2 状态空间模型
状态空间模型(SSMs)是一类通过潜在状态利用线性常微分方程框架,将一维输入转换为输出的工具。对于一个输入 和输出 的系统,其模型动态由以下方程描述:
其中,、、、和 是模型的参数。这些模型的离散版本,如 Mamba,采用零阶保持(Zero-Order Hold,ZOH)方法进行离散化,并允许模型通过选择性扫描机制适应输入数据。这种适应性在复杂应用中尤为重要,例如图像恢复任务,在这些任务中,理解不同图像区域之间的上下文关系是至关重要的。
3.3 模型架构
如图所示,该方法使用了一个多尺度的U-Net架构,并以三个分支的频率-状态空间模型(Frequency-SSM Blocks)作为核心组件。同时,我们还利用了Mamba模块的输入依赖性特性,通过不同尺度的降解图像生成注意力图,从而使模型能够捕捉不同尺度上的降解分布。
看了一下还是经典做法,先过一个3x3的卷积提取浅层特征,之后浅层特征再走一系列操作提取深层特征,之后两个特征相加再卷积输出重建后的图像。
3.3.1 FrepSSM块
FrepSSM块旨在通过结合三种不同分支的协同作用来应对复杂的挑战:傅里叶建模分支、空间分支和频带分支。
傅里叶建模分支
对于傅里叶建模分支,输入特征首先经过卷积层生成。通过快速傅里叶变换(FFT)转换为其傅里叶频谱,随后分解为幅度谱和相位谱。幅度谱和相位谱通过卷积块在频域中进行精细化处理,最后通过逆傅里叶变换(iFFT)返回到空间域:
其中,表示1×1卷积操作,表示一系列卷积操作和激活函数。此分支负责捕捉和处理图像的频域表示,完成图像的全局恢复。
与另一篇论文的傅里叶分支建模区别
博客:《FMambaIR:一种基于混合状态空间模型和频域的方法用于图像恢复》学习笔记-CSDN博客
一个走的卷积层,一个走的注意力机制
空间分支
我们首先使用层归一化处理输入特征 得到。然后,特征通过两个并行子分支。第一个子分支简单地使用SiLU激活函数。另一个子分支在1×1卷积后执行空间Mamba操作。空间Mamba包括以下序列:DWConv → SiLU → Mamba扫描 → LN,其中Mamba扫描指的是图所示的经典2D Mamba扫描。两个子分支的输出进行逐元素相乘,得到空间输出:
其中,表示逐元素乘法。
应该是引用这篇论文:MambaIR
博客:《MambaIR:一种基于状态空间模型的简单图像修复基线方法》学习笔记-CSDN博客
频带分支
经过二层小波包变换(Wavelet Packet Transform, WPT),以捕获不同频带的信息,输出大小为。这些频带信息从左上到右下重新排列回原始分辨率,随后通过频率Mamba进行处理,频率Mamba在图中有独特的频率扫描方式。由于图中的策略无法完全匹配低到高频的小波特征,我们采用类似于LocalMamba的策略。我们将小波特征分成四个块,并逐块进行扫描。
一张图像分成4个patch,每个patch 都是mamba的扫描方式。
处理后,经过对应的k层小波包逆变换,最后与逐元素相乘,得到频带扫描输出。这个过程可以表达为:
其中,WPT和IWT分别表示离散小波包变换和逆离散小波包变换。我们将Mamba建模的频域维度视为傅里叶建模与Mamba之间的过渡,实现了全局降解处理和局部细节恢复的无缝连接。
输入特征通过残差连接与空间分支的输出相加。最后,通过连接三个分支的输出特征并应用1×1卷积操作,块完成特征的协调合成。在图中,我们展示了不同分支的特征可视化,以说明各个分支在降雨捕捉上的差异。
(a) 雨天图像。(b) 注意力图。(c)、(d) 和 (e) 分别是空间分支、频带分支和傅里叶建模分支的特征图。空间分支 (c) 能全面识别雨滴,但边界较模糊。傅里叶分支 (e) 输出高对比度特征,更集中于较大的雨滴条纹。总体而言,频带分支 (d) 的表现介于两者之间。
3.3.2 数据依赖性降解先验注意力图
在降解图像的不同区域,恢复的难度差异显著,这受到降解分布和图像背景复杂度的影响。如图所示,Restormer执行的恢复在山脉等区域表现出巨大的误差,这些区域的特点是密集的雨点和复杂的背景。这一差距源于缺乏对降雨图像降解的明确建模。
为了解决这个问题,我们利用Mamba模块的独特数据依赖性特性,允许它根据特征的重要性动态地关注或忽视特定输入特征。我们提出了一种创新方法,通过使用Mamba模块生成降解先验,从而增强模型对图像中不同区域降解程度的识别和处理能力。该过程包括通过在多个尺度上对输入图像进行空间2D扫描来生成降解先验。这些降解先验随后与相应尺度的特征逐元素相乘,并与特征相加:
其中, 表示低分辨率的降雨图像, 表示相同尺度下的特征。通过结合不同尺度的输入特征,我们的模型能够有效地划分不同粒度的降解位置。
本质上,这种方法利用Mamba模块的内在适应性,使其能够动态分配注意力到输入图像的不同区域,从而使我们的模型能够更好地应对图像降解分布的复杂性。
3.4 损失函数
除了新的模块和降解先验外,我们还引入了新的损失函数来优化网络的训练过程,以在空间和频率域都取得良好的效果。损失函数由三个部分组成:空间域损失、相位谱损失和幅度谱损失。
在空间域中,我们使用L1损失来计算最终输出与真实图像之间的差异,以实现监督学习。同时,为了实现更好的全局信息重建,我们分别计算幅度谱和相位谱的L1损失,从而得到幅度和相位谱损失,表达式如下:
最终,损失函数的整体组成可以简洁地表示为:
其中, 和 被经验性地设置为0.05。
4、实验
4.1 数据集与实现
数据集:为了验证,我们在广泛使用的Rain100H、Rain100L、Test1200和Test2800数据集上训练和验证我们的模型。Rain100L从BSD200中选取,只有一种类型的雨条,包含200对训练图像和100对测试图像。Rain100H包含1800对训练图像和100对测试图像,涵盖五种雨条方向。Test1200有三组分别包含12,000张轻度、中度和重度雨天图像的训练集,每组4,000张图像。相同的三组1,200张图像用于测试。Test2800包含14,000对图像,其中12,600对用于训练,1,400对用于测试。
实现细节:我们的模型使用Pytorch框架实现,并在NVIDIA RTX 3090 GPU上执行。每层的块数影响模型参数量和去雨效果。在权衡后,我们将每层的块数设置为,这可以在合理的参数量下实现良好的性能。我们使用渐进式训练策略。以RAIN100L数据集为例,我们将总迭代次数设置为75,000,并设置图像大小为,相应的批次大小为。我们使用Adam优化器,默认参数,通过最小化损失函数LtotalL_{\text{total}}Ltotal来训练网络。初始学习率设置为,然后使用余弦退火策略逐渐衰减至。
4.2 对比实验
在本节中,我们将我们的方法与当前最先进的去雨方法进行比较:DerainNet、UMRL 、RESCAN、PreNet 、MSPFN、SPAIR 、MPRNet 、Restormer、Fourmer 、IR-SDE、MambaIR 和VMambaIR 。
定量比较:我们使用YCbCr颜色空间中的Y通道计算峰值信噪比(PSNR)和结构相似度指数(SSIM)得分。表1报告了在四个数据集上的性能评估。可以看出,我们的方法在所有基线算法中取得了最佳性能。
定性比较:为了展示我们提出的FreqMamba模型在图像去雨任务中表现出的增强的保真度和细节水平,我们比较了来自Rain100L和Rain100H数据集的挑战性降解图像的视觉质量当面对复杂或非常严重的雨条时,我们的方法几乎达到了完美的结果。与之前的方法相比,我们的FreqMamba在全局和局部恢复方面表现完美。
4.3 消融实验
我们在Rain100L数据集上对模型的不同组件进行了消融实验。
空间Mamba分支的探讨:为了验证空间Mamba分支的有效性,我们将Mamba模块替换为标准卷积层,其他部分保持不变。表中展示的结果表明,性能显著下降,因为标准卷积层的有限感受野导致其建模能力较差。
傅里叶分支的探讨:我们移除傅里叶分支,仅保留空间分支和频带分支。缺少全局建模能力使得模型性能下降了0.33 dB,这表明在傅里叶域中建模全局信息的重要性。这个结果与我们的初衷一致。
频带分支的探讨:我们移除频带分支,仅保留空间分支和傅里叶分支。模型性能下降了0.1 dB,这表明新增的频维度对于性能的贡献。
注意力图的探讨:我们的退化先验注意力策略基于不同尺度的雨天图像自适应地学习退化分布。为了验证其有效性,我们在表中进行了一系列消融实验。可以清楚地看到,带有注意力图的算法结果优于未使用退化先验的情况。
损失函数的探讨:频率损失旨在直接强调全局频率信息的优化。我们分别移除频率损失及其两个组件,并在表中检查其有效性。结果表明,移除频率损失会导致性能下降,进一步证明了其重要性。
4.4 在其他任务上的扩展
为了展示我们FreqMamba的潜力,我们将其扩展到低光图像增强和现实世界图像去雾任务。
低光图像增强的扩展:低光图像增强主要关注照亮场景的黑暗部分并去除放大的噪声。我们的三分支结构在这种场景中表现良好。我们采用LOL-V1和LOL-V2合成数据集来评估我们方法的性能。我们选择了几种低光图像增强方法进行比较:RetinexNet、KinD、ZeroDCE、KinD++、URetinex-Net 、FECNet和 SNR-Aware 。表展示了定量比较结果。
现实世界图像去雾的扩展:现实世界去雾旨在从现实世界的雾霾图像中恢复出清晰场景。对于这一任务,我们应用了两个数据集:Dense-Haze和NH-HAZE。我们将我们的方法与其他方法进行比较,包括DCP、DehazeNet、GridNet、MSBDN和AECR-Net。我们在表中展示了定量比较结果。