雨、雪等恶劣天气下采集的图片会发生严重退化,进而影响后续的语义分割、目标检测等^[1-2]计算机视觉算法的表现.现有多数基于深度学习的算法仅能去除单一天气,对于不同的退化无法直接迁移扩展.因此,多天气退化图像恢复凭借其仅需一次训练便可同时去除多种天气退化的优势而逐渐受到关注.例如,Li等^[3]利用多天气编码器和单输出解码器的结构首次提出多天气退化图像恢复网络All-in-One.Valanarasu等^[4]将任务查询向量引入视觉Transformer (Vision Transformer, ViT) 模型来得到干净背景图像.

图像生成模型具有强大的表征能力和丰富的解空间,自出现起便在图像生成领域表现出显著优势,如生成对抗网络 (generative adversarial network, GAN)^[5]、变分自编码器 (variational auto-encoder, VAE)^[6].其中,GAN利用对抗训练来提高视觉保真度,但存在训练不稳定、模式崩溃和图像伪影等问题;VAE通过最大化证据下界来优化数据的对数似然性,但平衡问题和变量崩溃现象限制其实际生成表现.近年来,去噪扩散概率模型 (denoising diffusion probability model, DDPM)^[7]在多个下游计算机视觉任务上表现出巨大优势.DDPM的主要原理为先连续添加高斯噪声来破坏原始数据,再进行反向采样来恢复图像.基于此,提出一种用于多天气退化图像恢复的自注意力扩散模型 (Transformer-based diffusion model for All-in-One weather-degraded image restoration, AWIR-TDM),以退化图像作为条件来引导反向采样.

目前,基于DDPM的计算机视觉算法^[7-8]在噪声估计网络部分均采用类似U-Net的卷积网络结构.然而,Peebles等^[9]的工作表明,采用ViT的类条件扩散模型可以取得相比于常用U-Net架构更好的图像生成表现.但是,ViT中自注意力 (self-attention, SA) 的计算负担与输入特征分辨率呈平方关系,无法直接应用于高分辨率的多天气图像恢复任务.对此,Wang等^[10]基于Swin Transformer提出用于多种图像恢复任务的统一型图像恢复算法Uformer,在分割的窗口中分别计算自注意力来减少计算负担;Zamir等^[11]利用转置自注意力 (transposed self-attention, TSA) 强化特征学习,即将原自注意力的表征从空间维度转移到通道维度,在去雨、去噪等任务上取得良好表现.然而,以上方法未依据自注意力的内积特点来进一步优化自注意力的表征,导致其训练和推理均需要大量时间.于DDPM而言,本身训练和采样均需要大量时间,如果在噪声估计网络部分直接使用计算负担巨大的ViT会加剧对硬件环境的要求.同时,文献[12]中表明自注意力和TSA具有接近的表现.因此,提出一种次空间转置自注意力噪声估计网络 (subspace transposed Transformer for noise estimation, NE-STT). NE-STT先将输入特征映射到次空间,再利用TSA构成次空间转置自注意力 (subspace transposed self-attention, STSA) 来强化特征学习并显著减少计算负担.同时,STSA不需要将输入特征分割为多个不重叠的窗口,而是直接从输入特征得到分辨率水平较低的查询向量Q和键向量K来进一步计算自注意力,相比Swin Transformer分割窗口更符合自注意力机制提取特征全局信息的初衷.同时,文献[11]和文献[13]中的工作表明,门控机制能够有效提高ViT中前馈网络 (feed-forward network, FFN) 的非线性表征能力从而提升视觉自注意力网络ViT的整体表现.基于此,提出双分组门控前馈网络 (dual grouped gated feed-forward network, DGGFFN),其在ViT中前馈网络FFN的深度可分离卷积后采用双分组门控机制来提高非线性表征能力.

综上所述,本文提出一种用于多种天气退化图像恢复的自注意力扩散模型AWIR-TDM,是去噪扩散概率模型DDPM和视觉自注意力网络ViT在多天气退化图像恢复任务上的有效探索.另外,提出NE-STT,其包括次空间转置自注意力STSA和双分组门控前馈网络DGGFFN.NE-STT一方面在次空间计算自注意力来减少计算负担,另一方面采用双分组门控机制来提高前馈网络的非线性表征能力,从而令噪声估计网络NE-STT估计出更加准确的噪声分布.相比常见的天气退化图像恢复算法,AWIR-TDM针对当前关注较少的多天气退化图像恢复任务,更加具有实际应用前景,并在去雪、去雨雾和去雨滴3个天气退化去除任务上取得超越现有算法的表现.

如图1所示,目前用于天气退化图像恢复的算法可以分为3类,即专一型、统一型和同一型.其中,专一型^{[14⇓⇓-17]}为不同任务设计专门的网络结构,模型架构与参数均不相同;统一型^[10-11]为不同任务设计统一的网络结构,模型架构一致但参数不相同;同一型^{[3-4,18 -19]}根据不同天气退化的相似特点设计同一的多天气退化图像恢复网络,在同时存在多种天气退化因素的数据集上训练, 对于不同天气退化图像恢复任务模型的架构和参数相同.

去噪扩散概率模型DDPM包括前向扩散过程和反向采样过程.前向扩散过程被定义为马尔可夫链,在连续的节点不断加入高斯噪声来获得有噪声的样本,进而逐步将高斯噪声分布转换为生成模型所训练的数据分布.具体而言,给定一个数据样本x₀~q(x₀),在多天气退化图像恢复任务中,x₀即为无退化的干净背景,q(x₀)为干净背景图像对应的原始分布,前向加噪过程定义为

式中:x_t为t时刻的加噪数据;T为加噪次数,t∈{0,1,…,T}.先验分布q(x_t|x_t-1)的数学形式为

式中:N为高斯分布;I为单位矩阵;β_t为噪声调节因子,其在前向扩散过程中线性地从 0.000 1 增大到0.02.则状态x_t可以用x_t-1表示为

式中:ε_t为t时刻的噪声分布.根据马尔可夫性,进一步得到x_t与x₀的关系为

式中:α_t=1-β_t,α-t=∏i=0tα_t.也可以表示为

式中:ε为标准高斯分布,即在前向扩散过程中每一步添加的噪声均是标准高斯噪声.前向扩散过程的结果是得到一个近乎服从高斯分布的数据,反向采样过程即学习一个参数化的后验分布p_θ(x_t-1|x_t)通过x_t生成前一个状态x_t-1.初始后验分布满足p(x_T)=N(x_T; 0, I),则反向采样过程定义为

式中:参数化后验分布p_θ(x_t-1|x_t)定义为

高斯分布中参数化估计均值μ_θ(x_t, x₀)和方差σt2的数学形式为

ε_θ(x_t, t)为噪声估计网络 (noise estimation network, NE-Net) 预测所得t时刻的噪声分布ε_t,其定义为ε_θ(x_t,t)=ρ_θ(x_t,t),ρ_θ为NE-Net的参数.网络训练的目的是令t时刻实际的噪声分布与预测的噪声分布误差最小,最终优化所得的参数可以定义为

随后,反向采样过程中t-1时刻的状态x_t-1可以用t时刻的状态x_t表示为

由此利用噪声估计网络来逐步反向采样出干净的背景图像.

如图2所示,条件扩散模型即在去噪扩散概率模型DDPM的反向采样过程中引入类别或图像来引导生成目标图像,不影响前向扩散过程.对于多天气退化图像恢复而言,将退化的图像作为条件引入采样过程,则反向采样过程变为

式中:x^即作为条件的天气退化图像.此时x_t-1~p_θ(x_t-1|x_t, x^),根据文献[8]中的条件反向采样过程为

ε_θ变为

即利用引入图像条件的噪声估计网络NE-STT来逐步反向采样得到去除天气退化的图像,如图3所示,图中f表示式(12).采用图像条件扩散模型中常用的通道维度拼接来引入图像条件.

在目前基于去噪扩散概率模型DDPM的图像恢复任务中,噪声估计网络NE-Net均采用以卷积神经网络CNN为核心的U-Net架构.文献[9]中表明视觉自注意力网络ViT的全局特性或架构设计相比CNN在建模特征依赖上更具优势.然而,ViT中自注意力巨大的计算负担和扩散模型本身更长的反向采样时间限制了其在多天气退化图像恢复任务AIR上的实际应用.为此,提出计算负担更小且噪声估计表现更好的NE-STT,如图4所示.

图中:Layer Norm为层归一化操作;gelu为激活函数;C、L分别为输入特征的维数和长度;r为次空间变换系数;Linear为线性层.

认为Swin Transformer这种分割窗口的操作提取的是局部的全局性信息,而非全局的全局性信息.因此,所提噪声估计网络NE-STT利用次空间变换系数将输入特征映射到特征次空间,再在次空间特征上表征自注意力.同时,为避免Swin Transformer由于窗口不重叠而必须不断进行窗口移位的现象,采用3个重叠卷积层对应Q、K和V来分别进行次空间映射 (subspace pProjection, SP).然而,实验发现空间维度的计算负担是通道维度的数倍.因此,NE-STT仅采用转置自注意力来强化特征学习,从而进一步减少计算负担和反向采样时长.次空间转置自注意力可以公式化为

式中:α为一个可学习的尺变参数.

表1是不同自注意力计算方式的计算量定量对比.其中,次空间变换系数r设为0.5,输入特征值得结构定义为 (64×64, 32),即L=64×64,C=32.如表1所示,TSA相比标准自注意力计算方式在计算量上下降3个数量级,而所提次空间转置自注意力相比TSA在乘法计算量上下降68.77%, 在加法计算量上下降69.85%.

此外,Zamir等^[11]和Chen等^[13]的研究表明,在视觉自注意力网络ViT的前馈网络FFN中采用门控机制可以进一步提升其非线性表达能力.受其启发,提出双分组门控前馈网络DGGFFN,其在FFN中深度可分离卷积层后利用双分组门控机制来强化特征学习.具体而言,DGGFFN将深度可分离卷积所得特征按照通道维度分为两组,再对两组特征分别利用另一组特征进行门控增强,最后将所得结果相加得到门控输出.双分组门控机制可以公式化为

式中:X₁、X₂表示将输入特征X^l分为两组;X^l、X^l+1分别为双分组门控机制的输入和输出.在估计噪声分布时,NE-STT以退化图像为条件并将其与前一次噪声状态按通道维度拼接作为输入,进而得到受退化图像条件引导的噪声分布.

本文算法在NVIDIA Tesla A100显卡上训练,在NVIDIA RTX 3090 显卡上分析.训练迭代次数为 100 万次,每次参与训练的图像对数量为16.网络参数优化器选择Adam,其学习率设置为常数1×10^-4,参数β₁=0.5,β₂=0.999.扩散步数为 2 000.在噪声估计网络NE-STT中,次空间变换系数在特征编码阶段分别为0.25、0.5、0.5、1,自注意力头的个数对应分别为1、2、4、8.

Allweather:根据文献[3]将Snow100K^[14]、Out-door Rain^[20]和Raindrop^[17]数据集组合而成,其包括18 069对图像用来训练多天气退化图像恢复算法.

Snow100K:来自文献[14],训练集包括5万对图像,测试集根据雪的大小分为Snow100K-L、Snow100K-M和Snow100K-S这3个数据集,其中分别包括16 611、16 588和16 801对图像用来评估算法的去雪性能.

Out-door Rain:雨雾共存的数据集^[20],训练集包括 9 000 对图像.采用其中名为Test1的子集作为测试集,图像对的数量为750.

Raindrop:包括861对训练图像和两个图像对数量分别为58和249的测试数据集.实验采用的测试集为Raindrop-A^[17].

此外,为验证算法恢复自然天气退化图像的表现,选用文献[14]中提供的雪数据集Snow,其中包括1 329幅自然雪图,同时还选用来自文献[21]的雨滴数据集Raindrop、雨雾数据集RainMist和雨痕数据集RainStreak,其分别包括67、13和185幅自然天气退化图像.

为了验证本文算法的有效性,实验对比本文算法AWIR-TDM与近期表现良好的统一型图像恢复方法Uformer^[10]、Restormer^[11]和多天气退化图像恢复方法All-in-One^[3]、TransWeather^[4]在数据集Allweather上经过一次训练后分别在雪、雨雾和雨滴测试数据集上的表现.选取常用的有监督图像质量评价指标峰值信噪比 (peak signal noise ratio, PSNR) 和结构相似性 (structural similarity, SSIM) 进行定量对比.峰值信噪比的计算公式如下:

式中:S为测试集的个数,S=1时即代表计算单个测试集结果的指标;M_i为第i个测试集中图像对的数量;v_max为图像的像素值的最大值,一般取1或255;g_mse(X_{i, j}, Y_{i, j}) 对应第i个测试集中第j对图像X_{i, j}和Y_{i, j}之间的像素值均方差.结构相似性的计算公式如下:

式中:μXi, j、μYi, j分别是图像X_{i, j}和Y_{i, j}像素值均值;σXi, j、σYi, j分别是图像X_{i, j}和Y_{i, j}像素值标准差;σXi, jYi, j是图像X_{i, j}和Y_{i, j}像素值的协方差;C₁=(K₁v_max)²,C₂=(K₂v_max)²,其中K₁=0.01, K₂=0.03.

不同算法的多天气图像恢复表现定量对比如表2所示,其中数据由原文复现或源代码所得.另外,表中加粗、双下划线和单下划线数据分别对应测试集指标最高、统一型恢复方法恢复算法效果排名第二和同一型恢复算法效果排名第二;空白表示不适用,下同.由表可见,所提AWIR-TDM在图像恢复质量上均高于对比方法.具体而言,本文算法相比 Restormer 和TransWeather,在Snow100K的测试集上分别平均提升1.87和3.39 dB,在雨雾共存测试集Test1上分别提升1.98和3.18 dB,在雨滴退化数据集Raindrop-A上分别提升1.23和 2.12 dB.值得注意的是,Uformer在雨雾共存的测试集Test1上表现最差,这说明现有的统一型图像恢复算法不能简单应用到多天气退化图像恢复任务.

不同算法所得恢复图像的视觉对比如图5~7所示,分别为退化图像经过去雨滴、去雨雾和去雪后所得的恢复图像.在图5中,相比其他算法,AWIR-TDM方法能够有效去除较大雨滴并得到与无退化图像接近的结果.如第1行对比图像中Uformer和All-in-One未能去除雨滴,Restormer和TransWeather所得恢复图像中仍存在明显痕迹,而本文算法能够彻底去除雨滴.在图6去雨雾视觉对比中,所有算法均能有效去除雨水痕迹,但相比本文算法,其他算法无法有效恢复图像的细节信息.如第1行对比图像中,本文算法能够得到清晰的指示牌字符,而其他算法所得字符已经无法辨认.特别地,发现Uformer处理雨雾图像时未能去除图像边缘的雨雾.在图7第一行对比图像中,本文算法能够有效去除上衣口袋上的雪痕迹,而Uformer、All-in-One和TransWeather的结果中出现明显残留.尽管Restormer的结果中没有明显的雪痕迹,但却得不到清晰的细节纹理,而本文算法得到的纹理更接近真实背景图像.

此外,实验还进一步将本文算法用于统一型天气退化图像恢复任务,并分别进行对应的训练与测试,实验所得定量对比结果如表3~5所示.结果表明,AWIR-TDM算法取得优于现有统一型图像恢复算法的表现.具体而言, 相比现有图像去雪算法,

本文算法在Snow100K的测试集上平均PSNR高出1.81~7.09 dB; 相比现有图像去雨雾算法,本文算法在Test1测试集上平均PSNR高出0.69~11.51 dB;相比现有去雨滴算法,本文算法在Raindrop-A测试集上平均PSNR高出0.79~4.82 dB.因此,本文算法不仅可以作为多天气退化图像恢复算法,还可以作为统一型天气退化图像恢复方法,并且在两种情况下均超越现有表现良好的方法.

为验证本文算法在自然天气退化图像恢复中的表现,实验选择自然度图像质量评估器 (naturalness image quality evaluator, NIQE)、空间光谱熵质量 (spatial-spectral entropy quality, SSEQ) 和神经图像评估 (neural image assessment, NIMA) 3个无监督指标来定量对比恢复图像的自然度.其中,NIQE 和 SSEQ 两指标的值越小表示图像自然度越高,NIMA 反之.表6为对比方法与本文算法在自然天气退化图像数据集上的定量对比结果.结果表明,在4个真实天气退化图像数据集上,本文算法AWIR-TDM均取得第一或者第二的表现,其中在RainMist和RainDrop数据集上,本文算法所得恢复图像的定量指标均为最高.在Snow和RainStreak数据集上,本文算法所得恢复图像的NIQE指标和SSEQ指标分别排名第二,而其他指标均为所有对比算法中最高.

此外,不同算法所得的自然天气退化图像恢复视觉对比如图8所示.其中,在处理第1~3幅雨滴图像时,AWIR-TDM能有效去除雨滴痕迹,同时得到更符合视觉感知特点的纹理特征.如第1行对比图像中,All-in-One所得结果中除天空部分残留的雨滴外其他雨滴被彻底去除,但树枝原有纹理被破坏,从而导致视觉效果较差.在处理第4~6幅雨图时,本文算法有效去除退化图像中的雨水痕迹,而其他算法所得结果中均存在明显的雨痕残留.如第5行对比图像中,Uformer、Restormer等算法均无法去除肩膀处的雨痕.第7、8幅退化图像为雪图,其中第7幅图中由于雪的运动出现大量雪痕,第8幅图中存在大小不一的雪花颗粒,本文算法均有效去除雪痕和雪花,而其他算法无法适应不同的雪特点.如最后一行对比图像中,All-in-One和TransWeather等算法无法同时去除多样的雪花颗粒.

为验证所提NE-STT各组件的有效性,利用实验对其进行消融分析.实验中,以去噪扩散概率模型DDPM中采用残差块 (residual block, ResBlock) 和空间自注意力的架构作为基线网络Baseline;随后分别对比采用ViT架构、TSA、STSA、FFN、单门控前馈网络 (single gated feed-forward network, SGFFN)和DGGFFN对最终恢复效果的影响以及模型的单步噪声估计用时.扩散模型采样耗时,因此实验选择图像数量较少的测试集Raindrop-A.表7为NE-STT各部分组件的消融实验结果.其中,单步估计用时为噪声估计网络NE-STT处理100张大小为6像素×64像素×64像素的图像输入时所需时间的平均值.

由表可见,噪声估计网络NE-STT中双分组门控前馈网络DGGFFN的使用较未加门控机制增强的噪声估计网络在单步估计用时上增加 0.011 9 s,但相比基线网络Baseline和其他网络结构耗时均明显减少.具体而言,AWIR-TDM单步估计用时分别为Baseline和ViT的61.50%和38.04%.同时,在Raindrop-A数据集上,AWIR-TDM相比Baseline和ViT在平均PSNR上分别提升3.16和 2.47 dB.STSA相比SA在平均PSNR上提升0.68 dB,相比TSA在平均PSNR上降低1.05 dB,但在单步估计用时上减少53.96%.随后,再通过DGGFFN进一步提升1.84 dB,从而相比采用TSA的ViT结构提升0.79 dB.实验表明,噪声估计网络NE-STT中包括的次空间转置自注意力机制和双分组门控前馈网络DGGFFN是有效的,并且显著减少噪声估计用时.

针对多天气退化图像恢复任务,本文提出一种有效的算法AWIR-TDM,利用去噪扩散概率模型DDPM和视觉自注意力网络ViT来条件获取高质量的恢复图像.具体而言,AWIR-TDM利用天气退化图像来条件引导DDPM的噪声估计,逐步反向采样得到干净背景图像.在噪声估计网络NE-Net中,提出一种更加快速且噪声估计更加准确的次空间转置自注意力噪声估计网络NE-STT.其中包括次空间转置自注意力STSA和双分组门控前馈网络DGGFFN.STSA利用次空间变换系数将输入特征映射到次分辨率空间,进而再利用计算量更低的转置自注意力TSA强化特征学习,其能有效学习输入特征的全局性依赖.DGGFFN在前馈网络FFN中的深度可分离卷积后采用双分组门控机制,相比无门控或单门控FFN,其能够进一步提高前馈网络的非线性表征能力.实验表明,AWIR-TDM所得恢复图像在定量指标和视觉效果上优于现有多天气退化图像恢复算法.然而,本文算法处理退化图像时缺乏实时性.因此,未来将致力于研究如何加速用于多天气退化图像恢复的扩散模型.