清华大学提出 TANet 用于综合恶劣天气图像复原的三重注意网络 !

科技   2024-11-17 17:25   广东  

不利天气图像修复旨在去除由不利天气条件(如雾、雨、雪等)引起的不 desired 退化图像。现有方法在处理单一天气条件时取得了显著成果。

然而,在遇到不可预测的天气条件时,它们面临挑战,这些条件在实际场景中经常发生。

虽然不同天气条件表现出不同的退化模式,但它们具有高度相关和互补的共同特性,例如退化模式引起的遮挡、颜色失真和大气现象引起的对比度衰减。因此,作者关注利用多个天气条件之间的共享知识以统一地恢复图像。

在本文中,作者提出了一种 Triplet Attention Network (TANet),以高效有效地解决一次性不利天气图像修复问题。

TANet 包括 Triplet Attention Block (TAB),该模块融合了三种注意力机制:局部像素级注意(LPA)和全局条纹注意(GSA)以解决非均匀退化模式引起的遮挡,以及全局分布注意(GDA)以解决大气现象引起的颜色失真和对比度衰减。

通过利用不同天气条件之间共享的通用知识,TANet 成功地以统一的方式解决多个天气条件。实验结果显示,TANet 在一次性不利天气图像修复方面高效地实现了最先进的结果。

源代码可在本站 https://github.com/xhuachris/TANet-ACCV-2024 中找到。

1 Introduction

恶劣的天气条件(如雾、雨、雪等)通常会导致图像中的 unwanted artifacts,从而降低图像的质量。这些退化模式会模糊图像的结构和细节,严重影响许多下游视觉任务。恶劣天气图像修复的目标是从单个退化图像中去除不必要的退化模式。然而,由于天气条件通常是非均匀的和时间变化的,恶劣天气图像修复是一个具有高度非凸问题的任务,已在计算机视觉领域广泛研究[12, 17, 25]。

恶劣天气图像修复随着深度学习的发展取得了显著的进步。在单一天气图像修复任务上,已有几项研究取得了有前途的结果,包括去雾,去雨 和去雪 。尽管以前的工作在特定天气条件下显著提高了视觉效果,但需要先验特定天气条件知识,这限制了它们在不可预测的现实生活中应用的可能性。由于天气条件会随时间变化,因此图像修复模型处理不可预测天气条件变得更加具有挑战性。因此,需要开发一个可以统一处理多种天气条件的一体化图像修复网络,而不需要依赖特定的天气先验知识。

最近,许多研究致力于以统一的方式解决多种不良天气条件。为了处理各种退化模式,一些方法通过设计特定于天气的组件来提取特定于天气的特征,例如天气类型 Query [33]和退化条件下的 Prompt [28]。然而,尽管不同类型的天气条件导致不同的退化模式,但这些退化模式高度相关且互补。

例如,雨雪 Mask 通常表现出各种方向和大小的 Mask 伪迹。大气颗粒的散射在不良天气条件下常常导致颜色失真和对比度衰减[36]。因此,作者不是设计特定于天气的模块来区分不同类型的天气条件,而是利用退化模式之间的共性来解决所有不良天气图像修复。这种方法使作者能够基于不良天气条件的归纳偏见构建一个有效且高效的形象修复网络。

在本文中,作者提出了一种名为Triplet Attention Network(TANet)的算法,该算法包括一个包含三种类型注意力模块的Triplet Attention Block(TAB),以有效且高效地解决一次成像的全天气图像恢复,包括去雾、去雨和去雪。如图1所示,雾天、雨天和雪天模式包含通常非均匀且大小不同的 Mask 伪影。为了有效处理这些退化的模式,TAB 引入了两种类型的空间注意力模块。

具体而言,TAB 使用局部像素级注意(LPA)来捕捉局部空间信息,以及全局条纹级注意(GSA),包括水平条纹注意和垂直条纹注意,来捕捉全局空间信息,这使得 TAB 能够处理具有各种方向和大小的雨天和雪天模式。此外,多尺度设计也使得 TAB 在各种恶劣天气条件下有效地处理非均匀退化模式。

此外,在恶劣天气条件下拍摄的图像通常会受到大气粒子散射的影响,导致颜色失真和对比度减弱。大气粒子的分布对散射强度有着重要影响。因此,TAB 采用了全局分布注意力(GDA)来捕捉大气粒子的变化分布。

具体来说,由于大气粒子的分布在不同退化的图像中有所不同,TAB 利用实例归一化在每个图像内部进行特征归一化,从而允许自适应调整退化图像中的特征分布。

因此,通过使用三种注意力机制,TAB 可以关注局部、全局和分布信息,利用恶劣天气条件的归纳偏差,从而实现高效有效的网络。实验结果表明,TANet 在合成和实际恶劣天气图像修复数据集上都取得了最先进的结果。

TANet 的贡献可以总结如下:

  1. 作者提出了一种Triplet Attention Network(TANet),这是一种高效且有效的集成图像修复网络,用于解决恶劣天气条件下的问题。

  2. TANet 利用 Tri-local Attention Block (TAB),利用不良天气条件的归纳偏差同时捕获局部、全局和分布信息,以消除退化模式。

  3. 实验结果表明,TANet在合成和真实世界的恶劣天气图像恢复数据集上都实现了最先进的结果。

2 Related Work

单次降质图像修复

随着深度学习的发展,单天气图像修复已经取得了显著的成果,包括图像去雾,去雨,以及去雪。对于图像去雾,一些研究通过提取与雾相关特征来提高去雾性能。

例如,Deng等人[9]提出了一种雾感表示蒸馏模块来提取雾感特征。Guo等人[11]提出了一种CNN与Transformer混合网络,其中传输感知位置嵌入用于解决雾状图案。对于图像去雨,一些研究[16, 21, 31]利用基于递归的网络逐步移除雨条纹。

Li等人[21]提出了一种利用膨胀卷积的递归网络,以增强感受野。Jiang等人[16]提出了一种多级金字塔架构,以细粒度到大粒度的方式重复移除雨条纹。对于图像去雪,以前的研究主要关注解决不同大小的雪状图案。Chen等人[4]提出了一种大小 Aware 和透明度 Aware 网络,用于去除雪和遮盖效果。

Zhang等人[45]提出了一种多级雪去除网络,该网络利用语义和几何指导以细粒度到大粒度的方式进行处理。

虽然这些针对特定天气条件的方法在特定天气条件下取得了有前景的结果,但由于特定天气架构的设计,其扩展到其他天气条件的能力仍然是一个问题。因此,一些工作提出了通用图像修复网络来解决多种降质模式。

2.0.1 Multiple Degradation Image Restoration

一些方法[2, 8, 43]选择设计针对降质特性的架构,而另一些方法[2, 8, 43]则开发出通用的模型,可以用于不同的任务,以处理各种类型的降质。Zamir等人[43]提出了一种多块架构,用于以交叉到细粒度的形式持续恢复降质图像。Chen等人[2]设计了一种非线性激活无关的网络,用于通用图像恢复,无需使用非线性激活函数。

Cui等人[8]提出了一种双域选择网络,包含空间和频率选择,以提取图像恢复的关键特征。尽管这些方法使用通用架构来解决各种类型的降质,但它们需要手动切换预训练模型来处理不同类型的降质。这在实际应用中并不适用,因为降质模式通常是不可预测的。

此外,尽管可以将这些通用图像恢复方法以全面的方式进行优化,但它们通常忽略了恶劣天气条件的不利影响,包括由降质模式引起的遮挡人工制品、颜色失真和由大气颗粒散射引起的对比度衰减,限制了它们在处理恶劣天气条件时的性能。因此,作者提出了一种利用恶劣天气条件的归纳偏差来一次性解决各种未知天气条件的方法,即TANet。

2.0.2 All-in-one Image Restoration

与单降和多降图像修复相比,一站式图像修复旨在通过统一的模型解决各种类型的退化,这提供了几个优势,例如强大的泛化能力和减少的存储要求。

为了以一站式的方式解决各种天气条件,一些研究将特定的天气模块集成到一站式恢复网络中,以适应处理各种未知降质。Chen等人[6]提出了一种教师-学生架构,通过从多个特定的天气教师模型中蒸馏知识。然而,学习多个特定的天气教师模型显著增加了训练期间的计算成本。

Valanarasu等人[33]提出在Transformer中利用可学习天气类型 Query ,Potlapalli等人[28]提出在Transformer中引入可学习降质信息 Prompt ,以实现一站式图像恢复。然而,除了为一站式图像恢复学习额外的特定天气参数外,作者注意到不同类型的天气条件高度相关且互补。

这些降质模式具有共同的特征,如由降质模式引起的遮挡,以及由于大气粒子散射导致的颜色失真和对比度衰减。

因此,TANet利用三组注意力模型同时解决遮挡和散射问题,充分利用恶劣天气条件的归纳偏见,有效地且高效地解决恶劣天气图像恢复的一站式问题。

3 Proposed Method

Overview

在本文中,作者提出了一种Triplet Attention Network (TANet),以一体化的方式解决恶劣天气图像恢复问题。与之前的方法不同,TANet利用了恶劣天气条件的归纳偏差,从各种退化类型的通用知识中提取特征。如图2所示,TANet是一个基于CNN的编码器-解码器网络,首先包含两个特征嵌入层(FEL),用于降维,每个FEL包括一个具有三个残差块的卷积层。接下来,作者堆叠多个Triplet Attention Block(TAB)。

每个TAB包括局部像素相关注意力(LPA)、全局条纹相关注意力(GSA)和全局分布注意力(GDA),同时解决遮挡和散射伪影。

最后,作者使用另外两个FEL将关注到的特征上采样,用于恢复干净的图像。后续部分将详细介绍TAB组件,包括LPA、GSA、GDA和用于优化TANet的最终损失函数。

Triplet Attention Block (TAB)

如图2所示,TAB包含局部像素注意力(LPA)、全局条纹注意力(GSA)和全局分布注意力(GDA),同时解决遮挡和散射伪影问题。

输入特征,其中分别表示高度、宽度和通道数。作者对通过卷积层处理,然后将其传递给三个并行的分支,以生成多尺度特征,其中分别是LPA和GSA的输出。


接下来,作者将这些特征连接起来,然后通过一个具有残差连接的卷积层生成多尺度关注的特征


在这一步中,作者在恶劣天气条件下处理非均匀降质模式。在解决掩盖目标结构的降质模式之后,作者利用GDA解决由大气颗粒散射引起的光彩失真和对比衰减问题。因此,作者通过GDA对进行处理,采用两个残差连接全局解决大气颗粒的未知分布问题。


以下是作者详细描述LPA、GSA和GDA的过程:

首先,作者来介绍LPA(Local Pattern Analysis)。LPA是一种自适应特征提取方法,它能够有效地捕捉输入数据中的局部模式。具体来说,LPA通过计算每个局部区域内的特征向量来表示输入数据,这些特征向量能够捕捉到输入数据中的局部结构信息。通过这种方式,LPA能够有效地降低计算复杂度,提高模型的效率。

接下来,作者介绍GSA(Gradient Shift Analysis)。GSA是一种基于梯度分析的方法,用于提取输入数据中的局部特征。具体来说,GSA通过对输入数据进行局部线性变换,使得输入数据的局部梯度发生改变。通过分析这些改变后的局部梯度,GSA能够有效地提取出输入数据中的局部特征。

最后,作者介绍GDA(Gradient Ascent Analysis)。GDA是一种基于梯度上升分析的方法,用于提取输入数据中的局部特征。具体来说,GDA通过对输入数据进行局部线性变换,使得输入数据的局部梯度发生改变。然后,作者通过梯度上升算法来优化这些局部线性变换,从而提取出输入数据中的局部特征。

通过以上详细描述,作者可以看出LPA、GSA和GDA都是用于提取输入数据中的局部特征的方法,它们各自具有不同的特点和优势。在实际应用中,作者可以根据具体问题的需求,选择合适的方法来提取输入数据中的局部特征。

3.2.1 Local Pixel-wise Attention (LPA)

如图2所示,为了有效解决由短程降质模式引起的重叠问题,TAB利用局部像素级注意(LPA)提取局部空间特征。受到[38]的启发,输入特征,作者利用在通道轴上进行平均池化和最大池化操作,生成两种类型的特征图:

接下来,将它们 ConCat ,并跟随一个带有符号函数的卷积层,生成空间注意力特征如下:


σ 表示 sigmoid 函数。

3.2.2 Global Strip-wise Attention (GSA)

如图3所示,为有效解决由长程退化模式引起的重叠问题,TAB利用GSA提取全局空间特征。受到[13]的启发,由于被遮挡的伪造物,如雨雪图案,包含具有各种方向的退化模式,作者利用包含水平条纹池化和垂直条纹池化操作的条纹池化,将特征投影到水平和垂直方向。

这使得作者可以有效地解决具有各种方向的长程退化模式。设输入特征为,作者通过长程平均池化操作将投影到水平和垂直特征。



其中, 是 GSA 的最终输出, 分别表示 sigmoid 函数和逐元素乘法。

3.2.2 Global Distribution Attention (GDA)

最后,如图2所示,为了解决大气颗粒散射引起的光彩失真和对比度衰减问题,TAB利用GDA捕获大气颗粒的分布。

为了适应地捕捉各种退化图像的特征分布,作者采用实例归一化,该方法在实例内部进行归一化。受[3]的启发,作者首先通过卷积层对输入张量在通道维度上进行划分。


接下来,作者通过实例归一化(IN)层处理,同时通过一个不使用实例归一化层的卷积层处理。这使得TAB能够自适应地调整特征分布,同时同时保留原始信息。


其中 是最终的输出。 分别表示均值和方差运算。 是可学习的仿射参数。

3.2.3 Loss Function

为了优化TANet,作者采用Charbonnier损失 作为损失函数。

在论文中,O和G分别表示恢复图像和真实图像,且

4 Experiments

Datasets and Implementation Details

作者利用三个合成数据集来优化TANet,包括RESIDE[18]用于去雾,Rain1400[10]用于去雨,以及Snow100K[24]用于去雪。具体而言,RESIDE包括包含110,000个室内训练配对的ITS数据集,以及包含313,950个室外训练配对的OTS数据集。Rain1400有12,600个训练配对,而Snow100K有100,000个训练配对。

作者遵循[6, 47]的统一采样策略,从每个数据集中各均匀采样5,000个训练配对,并将它们混合在一起形成一个训练集,包含15,000个训练配对。为了进行公平的比较,作者在这个混合训练集上优化TANet以及所有比较方法。对于评估,作者利用三个合成数据集和一个真实世界测试集来展示TANet的有效性。

对于合成测试集,作者选择包含500个室内配对和500个室外配对的SOTS[18]测试集用于去雾,包含1,400个配对的Rain1400测试集用于去雨,以及包含16,801个配对的Snow100K-L测试集用于去雪。对于真实世界测试集,作者选择包含4,500个配对的WeatherStream[44]测试集用于去雾,包含4,800个配对的Rain1400测试集用于去雨,以及包含3,960个配对的WeatherStream[44]测试集用于去雪。

4.1.2 Implementation Details.

作者在Pytorch平台上训练作者的方法,使用Adam优化器,初始学习率为,通过余弦退火策略逐步降低至

此外,作者采用数据增强,包括随机裁剪、翻转和旋转,其中作者随机裁剪图像至大小。作者用500k次迭代训练TANet,批量大小为16。作者使用NVIDIA RTX A5000 GPU训练和测试TANet。

Experimental Results

在本节中,作者对比了TANet与四种最先进的整体图像修复方法,包括Transweather[33],KCKE[6],WGWS[47],以及PromptIR[28],以及四种最先进的多种退化图像修复方法,包括NAFNet[2],FocalNet[8],GRL[22],以及MPRNet[43]。请注意,作者在第4.1节中描述的混合数据集上训练所有方法,并在不同的天气条件下分别评估它们。

4.2.1 Quantitative Comparisons.

在表格1中,作者在合成数据集上比较了TANet与最先进的算法,包括用于去雾的SOTS,用于除雨的Rain1400和用于去雪的Snow100K-L。在平均值上,TANet比最先进的MPRNet方法提高了1.26dB。尤其是,在去雾方面,TANet比MPRNet提高了3.01dB,在除雨方面提高了0.15dB,在去雪方面提高了0.62dB。

与最先进的一体化图像恢复方法PromptIR相比,TANet平均提高了2.18dB。在表格2中,作者在现实世界的WeatherStream数据集上比较了TANet与最先进的算法,提供了相应的真实雾天、雨天、雪天图像和对应的清晰图像。TANet也实现了最先进的结果,并比第二好的PromptIR提高了0.12dB。

此外,以前的图像恢复方法通常包含大量参数来解决多个未知天气条件,例如Transweather的38M个参数,KCKE的29M个参数和PromptIR的33M个参数。

相比之下,由于TANet利用了恶劣天气条件的归纳偏见,TANet只能使用9M个参数来解决所有一体化恶劣天气图像恢复的最先进结果。此外,TANet在NVIDIA RTX A5000 GPU上的推理时间也高效,为256x256大小的图像计算推理时间为18ms。

4.2.2 Qualitative Comparisons.

在图4至图6中,作者展示了在合成数据集上的定性比较,包括去雾(SOTS)、排水(Rain1400)和去雪(Snow100K-L)分别。结果表明,TANet相较于之前的最新方法实现了更好的或可比较的结果。在图7至图9中,作者展示了在实际WeatherStream数据集上的定性比较,包括去雾、去雨和去雪的结果分别。T

ANet同样相较于之前的最新方法实现了更好的或可比较的结果。特别是,实际场景往往涉及混合不良天气条件,如雨雾或雪雾模式。TANet仍能成功去除这些混合退化模式,证明了集成训练策略的有效性。

Ablation Studies

在本节中,作者对提出的Triplet Attention Block(TAB)进行组件分析。TAB包括三种类型的注意力模块,包括局部像素级注意力(LPA)、全局条纹级注意(GSA)和全局分布注意力(GDA)。作者在合成数据集上进行消融研究,包括用于去雾的SOTS测试集,用于去雨的Rain1400测试集,用于去雪的Snow100K-L测试集。

在表3中,第一行(Net1)作者通过用卷积层替换所提出的组件(包括LPA、GSA和GDA),构建了一个基准。在第二行(Net2),作者展示了LPA的有效性,它旨在捕捉由降质模式引起的位置遮挡伪影。通过使用LPA,作者平均可以获得0.22dB的改进, compared to Net1。在第三行(Net3),作者展示了GSA的有效性,它旨在提取由降质模式引起的全局遮挡伪影。通过使用GSA,作者可以显著提高性能,与Net2相比,可以提高1.15dB。作者总结了在TAB中使用GSA的成功原因如下。

由于在恶劣条件下,遮挡伪影是非均匀的,具有各种方向和大小,GSA将特征解耦成水平和垂直方向,并将其映射到有效应对各种方向的降质模式。此外,通过结合LPA和GSA,TAB可以利用多尺度特征来应对具有各种大小的不均匀降质模式。在第四行(NET4),作者展示了GDA的有效性,它旨在捕获大气颗粒的分布。与Net3相比,GDA可以进一步提高性能0.56dB。由于在恶劣天气条件下,图像通常受到大气颗粒的散射,TANet利用GDA成功地捕获了大气颗粒的分布,有效地提高了恢复结果的质量。

最后,作者在最后一行(Net5)中证明了使用FFT损失的有效性。遵循[8]中的方法,作者在频域中利用FFT损失来监督恢复图像,与Net4相比,可以进一步提高性能0.36dB。消融实验表明,TAB成功地利用了恶劣天气条件下的归纳偏差来解决降质模式,并考虑了遮挡和散射的概念。

5 Conclusion

在本文中,作者提出了一种名为TANet的新网络,用于一站式恶劣天气图像修复网络。

由于在恶劣天气条件下拍摄的图像通常受到降质模式和大气粒子散射的影响,导致遮挡、颜色失真和对比度减弱,因此TANet利用多个恶劣天气条件下的这些共同特性,有效地以一站式的方式修复退化的图像。

特别是,TANet利用包含局部像素级注意(LPA)、全局条纹注意(GSA)和全局分布注意(GDA)的Triplet注意力块(TAB)有效地解决了恶劣天气条件下遮挡和散射艺术的影响。

通过利用恶劣天气条件的归纳偏差,TANet有效地并有效地实现了在一站式恶劣天气图像修复方面的最先进性能。

参考文献

[0]. TANet: Triplet Attention Network for All-In-One Adverse Weather Image Restoration.

机器学习算法那些事
号主是大厂人工智能专家,专注于机器学习,深度学习以及计算机视觉等研究方向,每天会更新人工智能最前沿知识和分享自己的论文总结和学习笔记,让你系统化的学习每个知识点,每天进步一点点。
 最新文章