GridFormer: Residual Dense Transformer with Grid Structure for Image Restoration in Adverse Weather Conditions
Abstract
恶劣天气条件下的图像恢复是计算机视觉中的一个难点。在本文中,我们提出了一种新的基于Transformer的框架GridFormer,它可以作为恶劣天气条件下图像恢复的支柱。GridFormer采用剩余密集Transformer块在网格结构中进行设计,并介绍了两种核心设计。首先,它在Transformer层使用增强的注意机制。该机制包括采样器和紧凑自关注阶段以提高效率,以及局部增强阶段以加强局部信息。其次,我们引入了一个残余密集Transformer块(RDTB)作为最后的GridFormer层。这种设计进一步提高了网络从之前和当前的局部特征中学习有效特征的能力。GridFormer框架在恶劣天气条件下的五种不同图像恢复任务中实现了最先进的结果,包括图像脱雾、除雾、脱雾和除雾、下雪和多天气恢复。源代码和预训练模型将被发布。
1 Introduction
在雨、雾霾和雪等恶劣天气条件下捕捉高质量图像是一项具有挑战性的任务,因为在这些条件下会发生复杂的退化。这些问题包括色彩失真、模糊、噪音、低对比度和其他直接降低视觉质量的问题。此外,这种退化可能导致下游计算机视觉任务的困难,如物体识别和场景理解。
传统的恶劣天气条件下的图像恢复方法通常依赖于手工制作的先验,如平滑度和暗通道,采用线性变换。然而,由于较差的先验泛化,这些方法在处理复杂天气条件方面的能力有限。近年来,基于卷积神经网络(convolutional neural network,CNN)的方法被提出来处理图像的去噪问题。这些方法侧重于使用特定的架构设计来学习从天气退化图像到恢复图像的映射,例如残差学习、多尺度或多阶段网络、密集连接、GAN结构以及注意机制。然而,这些方法通常是为单一的特定任务而设计的,可能不适用于多天气恢复。
最近,出现了一种新的方法来解决统一架构中多天气恢复的挑战。开创性工作提出了一个多编码器和解码器网络,每个编码器专门处理一种类型的退化。利用神经结构搜索对网络进行优化。随后的工程借鉴了这种结构,以提高多天气恢复性能。例如TransWeather网络,该网络采用自关注进行多天气恢复。虽然TransWeather比特定任务编码器网络更有效,但其性能受到网络中不同尺度特征融合利用不足的限制。近年来,一些研究集中在设计通用骨干网,以利用网络中的多尺度特征来完成视觉任务。例如HRNet 和HRFormer采用多分辨率并行设计构建,学习高分辨率表示。RevCo采用了使用列(每列是一个子网络)的设计,目的是学习解纠缠表示。这些方法在人体姿态估计、语义分割、目标检测等方面都有很好的效果。然而,目前还没有专门设计的基于Transformer的方法来有效地利用这些特征来恢复恶劣天气条件下退化的图像。
图1 在恶劣天气条件下对图像进行比较的结果。(上)特定天气恢复和(下)多天气恢复任务的结果,显示了最先进的PSNR性能。
在本文中,我们提出了GridFormer,这是一种基于Transformer的网络,用于恶劣天气条件下的图像恢复。GridFormer使用嵌入在网格结构中的残余密集Transformer块(RDTB)来挖掘分层图像特征。RDTB是GridFormer的关键单元,它包含具有密集连接的紧凑增强Transformer层,以及与局部跳过连接的局部特征融合。紧凑型增强Transformer层采用采样器和紧凑型自关注器来提高效率,并采用局部增强阶段来加强局部细节。我们在天气退化基准上评估GridFormer,包括RainDrop、SOTS-indoor、Haze4K 、Outdoor-Rain,和Snow100K,见图1。
总之,这项工作的贡献有三个方面:
- 统一框架:我们提出了一种新的统一框架GridFormer,它是专门为恶劣天气条件下的图像恢复量身定制的。这个创新的框架将残余密集Transformer块(RDTB)与网格结构无缝集成,创建了一个全面的架构。值得注意的是,在网格结构中结合RDTB使GridFormer能够有效地捕获分层图像特征。网格结构有利于不同空间尺度背景信息的整合,增强了网络对图像的有效还原能力。
- 紧凑增强的自关注:GridFormer引入了紧凑增强的自关注机制,这是一个重要的贡献。该机制增强了Transformer单元的局部建模能力,使Gridformer能够在恶劣天气条件下捕获细粒度细节,同时提高网络效率。
- 最先进的性能:我们通过将GridFormer应用于恶劣天气条件下的五种不同图像恢复任务,包括图像脱除、图像去雾、图像脱除和去雾、下雪和多天气恢复,展示了GridFormer的一般适用性。我们的GridFormer在特定天气和多天气恢复任务上都实现了最先进的新技术。
3 Method
3.1 动机和架构
3.1.1 动机
我们的动机是迫切需要恢复在不利天气条件下拍摄的图像的技术。与天气相关的因素,如雾霾、雨和雪,会显著影响图像的质量和感知,进而影响监视、自动驾驶和户外摄影等各种实际应用。开发GridFormer的主要目的是解决恶劣天气条件对图像质量造成的持续挑战。我们的目标是创建一个图像恢复框架,有效地处理一系列恶劣天气情况,从而提高受这些条件影响的图像质量。
3.1.2 框架
图2 GridFormer架构。它由网格头部、网格融合模块和网格尾部组成。首先将金字塔梯度图像送入网格头部,提取分层初始特征
。网格融合模块对初始特征进行进一步细化,生成特征
。最后,网格尾部重建清晰图像
。
如图2所示,GridFormer包含从天气金字塔梯度图像到恢复图像的三条路径,每条路径在不同的图像分辨率下进行恢复。在GridFormer中,高分辨率路径与网络中的低分辨率路径持续动态交互,以准确地去除天气影响,而低分辨率路径由于具有更大的接受域而提供有用的全局信息。每条路径由七个GridFormer层组成。不同的路径通过下采样层、上采样层和加权关注融合单元相互连接,组成GridFormer的列。由于采用三行七列的网格结构,可以有效地共享不同分辨率的信息。具体来说,GridFormer由网格头(GH)、网格融合模块(GFM)和网格尾(GT)三部分组成。我们将在下面介绍每个部分的详细信息。
网格头Grid Head
为了提取初始的多分辨率特征,我们使用网格头架构并行处理金字塔输入图像。网格头中的每条路径由一个特征嵌入层(通过3 × 3卷积实现)和一个GridFormer层组成。如图2所示,给定天气退化图像,网格头从金字塔图像
(
和
分别为1/2和1/4尺度)中提取不同通道(即C,2C和4C)的层次特征
。在我们的实验中,我们使用C =48。网格头计算可定义为:
其中 i 为第 i 条网络路径,为特征嵌入层。符号↓表示下采样层,在该层中,我们使用3x3卷积和像素解洗刷操作将空间维度上的特征减半,同时将通道加倍。GFL是一个GridFormer层,主要由剩余的密集Transformer块构建而成。
网格融合模块Grid Fusion Module
为了充分融合网络中不同行、不同列的层次特征,我们在网格头部和网格尾部之间设计了网格融合模块。提出的网格融合模块的结构被组织成一个二维网格模式。如图2所示,融合模块被设计成三行五列的网格状结构。特别是,每行包含五个连续的GridFormer层,保持特征维度不变。在列轴上,根据在网格中的位置,我们采用下采样层或上采样层来改变特征映射的大小进行特征融合。
图3 网格单元结构与信息流。(a)单个网格单元的结构由下采样层、GridFormer层、上采样层和注意力融合操作四部分组成。(b)融合模块中电网单元的信息流。
图3a显示了融合模块中具有代表性的网格单元。GridFormer层是由三个残余密集Transformer层(RDTL)和一个1 × 1卷积组成的密集结构,这将在下一小节中讨论。下采样层和上采样层是对称的,并使用3 × 3卷积与像素shuffle或像素unshuffle操作来改变特征维度。此外,考虑到不同尺度的特征可能不是同等重要的,我们使用一种简单的加权关注融合策略来实现来自不同行和列维度的特征融合。我们首先为不同的特征生成两个可训练的权重,其中每个参数是一个n维向量(n是特征的通道)。我们将这些加权特征相加,得到融合特征。网格融合模块中的网格单元为特征融合提供了不同的信息流,如图3b所示,这些信息流结合不同的互补信息引导网络产生更好的恢复结果。
网格尾部Grid Tail
为了进一步提高恢复图像的质量,我们设计了一个网格尾部模块来预测多尺度输出。网格尾部的结构与网格头部的结构是对称的。具体来说,每条路径由GridFormer层、3 × 3卷积和用于图像重建的长跳过连接组成。采用跳跃式连接将输入信息直接传输到网格尾部模块,保持了原始图像的颜色和细节。完整过程表述为:
其中,为GridFormer在第i条路径上的最终结果,
为3 ×3卷积,
,i∈{0,1,2}为网格融合模块的输出特征。为了优化网络参数,我们使用两种损失的组合来训练GridFormer,多尺度Charbonnier损失和感知损失,其中感知损失权重设置为0.1。接下来,我们详细介绍了用于构建GridFormer元素层的核心组件剩余密集Transformer块。
3.2 残差密集Transformer模块
先前工作表明,使用密集连接具有许多优点,可以缓解梯度消失问题,鼓励特征重用并增强信息传播。因此,我们建议设计密集连接的Transformer来构建GridFormer的基本层。具体来说,我们提出了使用不同设置的剩余密集变Transformer(RDTB)来组成GridFormer。
图4 提出了剩余密集Transformer块(RDTB)的结构。它包括三个残差密集Transformer层,一个用于局部特征融合的1 ×1卷积和一个用于局部残差学习的局部跳过连接。剩余密集Transformer层主要由所提出的紧凑增强Transformer层构成,其中包含紧凑增强自关注和FFN.
如图4所示,RDTB包含密集连接的Transformer层、局部特征融合和局部残差学习。在实现密集连接时,我们主要合并了三层剩余密集变Transformer(RDTL),其增长率设为16。这意味着每个单独的RDTL生成16个新的特性映射。这些新生成的特征映射随后与从前一层接收到的特征映射连接起来。在每个RDTL中,我们使用几个紧凑增强的Transformer层(CETL)和ReLU激活函数来提取特征,并采用1×1卷积来确保输入和输出特征的通道数量相同。对于局部特征融合和局部残差学习,我们在RDTB中引入1 × 1卷积和局部跳跃连接来控制最终输出。
图5 所提出的紧凑增强型Transformer层的示意图,由紧凑增强型注意力和前馈网络(FFN)组成。左:紧凑增强关注层,包含特征采样、紧凑自关注和局部增强三个步骤。H、W、c分别表示特征通道的高度、宽度和数量。R是特征采样率。©和⊕分别指连接和元素求和操作.
Transformer的直接应用对我们的网格网络将导致高计算开销,因此我们开发了一种具有成本效益的紧凑型增强关注,采用采样器和紧凑型自关注阶段来提高效率,以及局部增强阶段来增强Transformer中的局部信息。图5说明了提议的紧凑型增强注意力的详细结构。
特征采样Feature Sampling
我们首先设计了一个采样器来产生下采样的输入令牌,用于随后的自注意计算。采样层通过步长为r的平均采样层来构建样本。采样层不仅增加了接收域以观察到更多的信息,而且增强了对输入令牌的不变性。此外,生成的低分辨率特征可以减少后续层的计算。特征采样步骤表示为:
其中表示输入令牌。
为输出令牌。
表示步长为r的平均池化操作。在实验中,我们经验地将三行GridFormer层中的r分别设置为4,2和2。
紧凑型自监督Compact Self-attention
给定尺寸为H × W ×C的特征,最近基于低级别Transformer的方法,旨在探索键和查询之间的远程依赖关系来计算N × N注意力图(N = H ×W),这导致复杂性高,无法从渠道维度对全局信息进行建模。因此,为了更有效地计算自我注意,我们采用了一种不同的策略。具体来说,如图5所示,对于采样器的输出特征,我们首先进行分割操作,将其沿着通道维度进行分割,得到
和
。然后,我们对
和
应用一个具有重塑操作的卷积层,它将
和
分别投影到查询(
),键(
)和值(
)中。受现有方法的启发,我们交换了他们产生的值来执行多头自注意,可以改善
和
之间的交互性。与交叉关注中交换查询进行特征交互的方法相比,我们的方法交换了交互和特征融合的值,发现它有利于更好的恢复性能。最后,我们通过将两个多头自关注的输出串接并改变它们的维数,得到了结果。建议的紧凑型自我关注机制可表述为:
其中[ · ]表示连接操作。Transformer的主要计算开销主要来自自关注层。与最近采用空间建模的基于变换的方法相比,键查询点积交互的复杂性随着输入的空间分辨率(即O(N × N))呈二次增长。我们提出的紧凑型自注意通过跨通道而不是空间维度执行SA来解决这个问题,导致跨通道的交叉协方差计算产生隐式编码全局上下文的注意图。因此,我们紧凑的自我注意生成了一个大小为的注意图,而不是大小为
的巨大的常规注意图。因此,我们紧凑的自我关注成功地降低了复杂性。
局部增强Local Enhancement
如图5所示,我们在紧凑自关注尾部增加了一个局部特征增强阶段。这一阶段包括反卷积操作,有时被称为“转置卷积”,其中反卷积用于局部特征传播,1 × 1卷积用于局部融合:
其中为最终输出。
和deconv 是1×1卷积层和反卷积层。
3.3 损失函数
受已有研究成果的启发,结合Charbonnier损失和感知损失,使用联合损失函数来训练我们的GridFormer。我们将Charbonnier损失视为像素级损失,用于每个尺度下的恢复图像和地面真实图像之间,而感知损失用于帮助我们的模型产生视觉上令人愉悦的结果。Charbonnier损失定义为:
其中,和
分别为恢复后的图像和Ground truth图像,k为GridFormer中图像尺度等级的指数。常数ε被经验地设定为
。对于感知损失,继之前的工作之后,我们采用预先训练好的VGG19 ,从VGG19的Conv5_4层提取感知特征,然后使用L1损失函数计算恢复图像的感知特征与其对应的Ground truth之间的差值。这种有效的感知损失侧重于捕获高级语义信息,从而产生更清晰的边缘和视觉上吸引人的结果,同时确保计算效率。具体来说,感知损失如下:
其中,C、H、W为预训练VGGNet φ的Conv5_4层得到的特征映射的维数。
最终训练GridFormer的损失函数L如下所示:
其中表示Charbonnier损失,
表示感知损失。α是用来平衡这两种损失的超参数。在我们的实验中,它被经验地设置为0.1。
3.4 与现有方法的区别
HRNet ,HRFormer和RevCol使用类似网格的结构,它们与我们的GridFormer不同。首先,GridFormer直接从像素级捕获多尺度特征,而HRNet和HRFormer在特征层级别执行多尺度特征提取,而RevCol则不包含多尺度机制。其次,GridFormer集成了一种新的自关注机制,更有效地增强了多尺度特征的融合。这种方法将其与HRNet、HRFormer和RevCol区分开来,后者在其特征融合过程中不使用紧凑的自关注。第三,我们的网络是为恶劣天气条件下的图像恢复而精心设计的,力求产生高质量的图像。与HRNet、HRFormer和RevCol不同,它们不是专门为这个挑战设计的,我们的网络架构是唯一适合解决这个任务中固有的复杂性的。