Wavelet-based Fourier Information Interaction with Frequency Diffusion Adjustment for Underwater Image Restoration
- 摘要
- 介绍
- 相关工作
- 论文方法
-
- 整体架构
- 离散小波变换与傅里叶变换
- 频率初步增强
-
- [Wide Transformer Block](#Wide Transformer Block)
- [Spatial-Frequency Fusion Block](#Spatial-Frequency Fusion Block)
- 损失函数
- [Cross-Frequency Conditioner](#Cross-Frequency Conditioner)
- [Frequency Diffusion Adjustment](#Frequency Diffusion Adjustment)
摘要
水下图像受到复杂多样的退化,不可避免地影响水下视觉任务的有效性。然而,大多数方法主要是在图像的原始像素空间中运行,这限制了对水下图像频率特性的探索,导致深度模型在产生高质量图像时没有充分利用其表征能力。在本文中,我们介绍了一种新的水下图像增强(UIE)框架,命名为WF-Diff,旨在充分利用频域信息和扩散模型的特性 。WF-Diff由两个可分离的网络组成:基于小波的傅里叶信息交互网络(WFI2-net)和频率残差扩散调整模块(FR-DAM)。通过对频域信息的充分探索,WFI2-net旨在初步实现小波空间中频率信息的增强。我们提出的FRDAM可以进一步细化初始增强图像的高频和低频信息,可以将其视为一个即插即用的通用模块,用于调整水下图像的细节。通过上述技术,我们的算法可以在真实的水下图像数据集上显示SOTA性能,并在视觉质量上达到具有竞争力的性能。代码可https://github.com/zhihefang/WF-Diff上获得。
介绍
水下图像修复是水下视觉领域中的一项实际而具有挑战性的技术,广泛应用于水下机器人[26]和水下物体追踪[6]等任务。由于水下场景中的光线折射、吸收和散射,水下图像通常受到严重的失真,表现为低对比度和模糊[2]。因此,清晰的水下图像在需要与水下环境互动的领域中起着至关重要的作用。水下图像增强(UIE)的主要目标是通过去除散射和校正退化图像中的颜色失真,从而获得高质量的图像。UIE对于与视觉相关的水下任务至关重要。
为了应对这一问题,基于水下图像物理特性的传统UIE方法被提出[15, 17, 29--31]。这些方法研究了由颜色偏差或散射引起的退化的物理机制,并通过补偿这些因素来增强水下图像。然而,这些基于物理学的模型由于表征能力有限,无法处理水下场景中所有复杂的物理和光学因素,从而导致在高度复杂和多样的水下场景中增强效果较差。近年来,一些基于学习的方法[7, 18, 28, 36]因其神经网络强大的特征表示和非线性映射能力,能够产生更好的结果。这些方法可以通过大量的配对训练数据学习图像从退化到清晰的映射。然而,大多数之前的方法都是基于图像的原始像素空间,未能充分探索水下图像在频率空间中的特性,这导致无法有效地利用深度模型的表示能力来生成高质量的图像。
【图1. 我们的动机。幅度和相位是通过快速傅里叶变换(FFT)得到的,重新组合的图像是通过逆傅里叶变换(IFFT)获得的。我们进一步探索了水下图像在小波空间中的频率特性。】
基于之前傅里叶方法的启示[12, 48],我们探索了傅里叶频率信息在UIE任务中的特性 ,如图1所示。给定两张图像(一个水下图像及其对应的地面真值图像),我们交换它们的幅度分量 ,并将其与对应的相位分量在傅里叶空间中结合 。重新组合的结果表明,视觉效果随着幅度交换而发生变化,这表明水下图像的退化信息主要包含在幅度分量中 。我们进一步探索了幅度分量在小波空间中的特性 。具体来说,图像可以通过离散小波变换(DWT)分解为低频子图像和高频子图像,然后我们交换低频子图像的幅度分量 。从视觉结果来看,我们发现了类似的现象,这意味着颜色退化信息主要包含在低频子图像中,而纹理和细节退化信息主要包含在高频子图像中。表1展示了不同频域策略的定量评估,证明了我们的发现是客观的。因此,如何充分利用频域信息的特性,并将其有效地融入到一个统一的图像增强网络中,是一个关键问题。
近年来,基于扩散 的方法[10, 35]因其在图像合成[23, 24, 32, 34, 52]和修复任务[5, 40, 46, 51]中的优异表现而受到广泛关注。这些方法依赖于分层去噪自编码器架构,使得它们能够迭代地逆转扩散过程,从随机采样的高斯噪声生成目标图像或潜在分布,并实现高质量的映射[10]。Tang等人[36]提出了一种用于水下场景的扩散模型图像增强方法。尽管标准扩散模型表现出足够的能力,但由于从随机生成的高斯噪声到图像的采样过程中引入的多样性,可能会出现一些预料之外的伪影 [45]。此外,扩散模型需要同时恢复图像的高频和低频信息,这限制了它们对细粒度信息的关注,导致错过了纹理和细节。因此,充分利用扩散模型强大的表示能力变得尤为关键。
在本文中,我们开发了一种新的水下图像增强(UIE)框架,充分利用频域信息和扩散模型的特性 ,命名为 WF-Diff ,该框架主要由两个阶段 组成:频率初步增强和频率扩散调整 。第一个阶段旨在利用频域特性,初步增强水下图像的高频和低频分量 。具体来说,我们首先通过离散小波变换(DWT)将输入图像转换到小波空间,得到一个代表输入图像低频内容信息的平均系数,以及三个代表输入图像稀疏的垂直、水平和对角线细节的高频系数。然后,我们设计了一个基于小波的傅里叶信息交互网络(WFI2-net),充分融合了Transformer [22] 和傅里叶先验信息的特性,分别增强高频和低频内容 。此外,为了实现高频和低频信息的交互,我们提出了一个跨频条件器(CFC),进一步提高生成质量 。第二阶段的目标是通过扩散模型对初步增强的粗略结果在细节和纹理上进行调整 。因此,我们提出了一个频率残差扩散调整模块(FRDAM) 。与之前的基于扩散的工作不同,FRDAM通过在小波空间中使用两个扩散模型,学习地面真值和初步增强结果之间的高频和低频信息残差分布,既能增强模型对细粒度信息的关注,又能减轻采样过程中多样性带来的不良影响。
总之,我们方法的主要贡献如下:
- 我们深入探索了水下图像在频域中的特性。基于这些特性和扩散模型,我们提出了一个新的UIE框架------WF-Diff,旨在实现频率增强和扩散调整。
- 我们提出了一个频率残差扩散调整模块(FRDAM),进一步优化初步增强图像的高频和低频信息。FRDAM可以视为一个即插即用的通用模块,用于调整水下图像的细节。
- 我们提出了一个跨频条件器(CFC),以实现高频和低频信息的跨频交互。
- 与现有的最先进方法(SOTA)相比,实验结果显著表明,我们提出的WF-Diff在性能上优于之前的UIE方法,并且通过广泛的消融实验验证了我们贡献的有效性。
相关工作
水下图像增强
目前,现有的水下图像去雾(UID)方法可以简要地分为物理模型和深度模型两种方法[15, 18, 28--30, 36]。大多数基于物理模型的UID方法利用先验知识来建立模型,如水下暗通道先验[29]、衰减曲线先验[38]、模糊先验[4]等。此外,Akkaynak和Treibitz[1]提出了一种基于修订物理成像模型的方法。然而,水下场景的深度图难以获取,这导致了不稳定的性能,通常受到严重的颜色偏差和伪影的影响。因此,手动建立的先验模型限制了模型在复杂多变的环境下的鲁棒性和可扩展性。
近年来,基于深度学习的方法[18, 28, 36]已取得了较为可接受的性能。为了缓解对真实水下配对训练数据的需求,许多方法引入了基于GAN的框架进行水下图像增强(UIE)[7, 14, 21, 49],例如WaterGAN [21]、UGAN [7]和UIE-DAL [37]。最近,一些复杂的框架被提出并取得了最先进的性能[15, 29]。例如,Ucolor[19]结合了水下物理成像模型,并设计了一个中介传输引导模型。Yang等人[43]提出了一种反射光感知的多尺度渐进修复网络,用于在各种水下场景中获得颜色均衡和丰富纹理的图像。Huang等人[13]提出了一种基于均值教师的半监督网络,能够有效地利用未标注数据中的知识。然而,大多数之前的方法都是基于空间域,未能充分探索水下图像的频率空间,这导致无法有效地利用深度模型的表示能力。
扩散模型
最近,扩散概率模型(DPMs)[10, 35]已广泛应用于条件图像生成[5, 40, 42, 46, 50]。Saharia等人[33]提出了Palette,展示了扩散模型在条件图像生成领域(包括上色、图像修补和JPEG修复)中的优异性能。Tang等人[36]提出了一种基于扩散模型的水下图像增强方法。然而,扩散模型的反向过程从随机采样的高斯噪声开始,生成完整的图像[45],由于采样过程的多样性,这可能导致意外的伪影。此外,扩散模型需要恢复图像中的高频和低频信息,这限制了其在关注细粒度信息方面的能力。因此,如何将扩散模型融入一个统一的水下图像增强网络是一个至关重要的问题。
论文方法
整体架构
给定一张水下图像作为输入,我们的目标是学习一个网络,生成一个输出,去除输入图像的颜色偏差,同时增强图像的细节 。WF-Diff的整体框架如图2所示。WF-Diff旨在充分利用频域信息的特性和扩散模型的强大能力 。具体来说,WF-Diff由两个可拆卸的网络组成 :基于小波的傅里叶信息交互网络(WFI2-net)和频率残差扩散调整模块(FRDAM)。我们首先通过离散小波变换(DWT)将输入图像转换到小波空间,得到一个低频系数和三个高频系数 。WFI2-net致力于实现频率信息的初步增强 。我们充分整合Transformer和傅里叶先验信息的特性,设计了宽Transformer块(WTB)和空间-频率融合块(SFFB),分别增强高频和低频内容。FRDAM由低频扩散分支(LDFB)和高频扩散分支(HDFB)组成,旨在进一步调整初步增强图像的高频和低频信息 。需要注意的是,我们提出的FRDAM通过使用两个扩散模型分别学习地面真值和初步增强结果之间的高频和低频信息的残差分布 。此外,所提出的跨频条件器(CFC)旨在实现高频和低频信息之间的跨频交互。
离散小波变换与傅里叶变换
离散小波变换(DWT)已广泛应用于低层视觉任务[11, 16]。我们首先使用DWT将输入图像分解为多个频率子带,以便分别实现低频信息的颜色校正和高频信息的细节增强 。给定一张水下图像作为输入
,我们使用Haar小波对输入进行分解 。Haar小波由低通滤波器 L 和高通滤波器 H 组成,如下所示:
我们可以获得四个子带,其可以表示为:
其中,
分别表示输入图像的低频分量和垂直、水平、对角方向的高频分量 。更具体地说,低频分量包含了输入图像的内容和颜色信息,而其他三个高频系数则包含了全局结构和纹理的细节信息 [31]。这些子带的分辨率是输入图像的一半,但由于DWT的双正交性质,它们不会导致信息丢失 。对于低频分量 ILL,我们将在傅里叶空间中进一步探讨其特性。
接着,我们引入傅里叶变换的操作[48]。给定一个图像
,其形状为 H x W,傅里叶变换 F 将 x 转换到傅里叶空间 X ,其表达式为:
其中, R(x) 和 I(x) 分别表示 X(u, v) 的实部和虚部。需要注意的是,傅里叶操作可以在每个通道的特征图上独立计算。根据图1和表1(我们的动机),我们得出结论:水下图像的颜色退化信息主要包含在低频子带的幅度分量中,而纹理和细节退化信息则主要包含在高频子带中。
频率初步增强
基于上述分析,在频率初步增强阶段,我们设计了一个简单但有效的WFI2-net,采用并行的编码器-解码器(类似U-Net)的结构,分别恢复低频信息的幅度分量和高频分量 。我们还利用跳跃连接将编码器和解码器中相同层次的特征连接起来。对于高频分支,我们利用变压器(Transformer)建模全局信息的优势来增强高频系数。我们设计了宽变压器块(WTB),使用多尺度信息,旨在建模长距离依赖关系 。我们的低频分支旨在恢复傅里叶空间中的幅度分量 。为了获得丰富的频率和空间信息,我们设计了空间-频率融合块(SFFB)。
Wide Transformer Block
Spatial-Frequency Fusion Block
损失函数
Cross-Frequency Conditioner
类似地,可以获得高频嵌入的VT和低频嵌入的VF:
输出特征图Tout和Fout然后可以从以下公式获得:
其中R表示复制操作,√ dk是矩阵Q的列数。