Trident去雾网络
【Trident:三齿的,三叉戟】
摘要
针对现有的去雾方法对非均匀雾霾的鲁棒性差,以及高雾霾区域的信息未知且难以估计,导致去雾效果模糊的问题,提出了一种由粗到精的模型Trident Dehazing Network(TDN),用于学习从有雾霾到无雾霾的图像映射,实现雾霾浓度的自动识别。具体地说,TDN由三个子网组成:编码器-解码器网(EDN)是TDN的主网,用于重建粗的无模糊特征;细节细化子网(DRN)用于细化编码器池化层中容易丢失的高频细节;而霾密度图生成子网络(HDMGN)能自动区分厚霾区和薄霾区,以防止在不同雾度密度的区域中的过度去雾或去雾不足。此外,我们提出了一个频域损失函数,使不同频段的监控更加统一。在合成和真实的数据集上的大量实验结果表明,我们提出的TDN优于现有的TDN。该方法具有较好的逼真度和感知能力,对高浓度雾霾和非均匀雾霾场景都有较好的推广效果,在NTIRE 2020非均匀雾霾去雾挑战赛中获得第一名。
1.介绍
雾是图像退化的主要来源之一,其导致低对比度、颜色失真和模糊问题。先前的工作[21,24]已经确定雾严重影响高级任务的性能,例如分类或语义分割,因此非常需要恢复所描绘的主体的预处理。为了去除雾并提高雾图像的可见性,已经提出了许多去雾方法[17,37,29,10,26,11,20,14,21,23,27,35]。
在早期基于大气散射模型的去雾方法中,如暗通道先验法[17]、对比色线[15]和雾线先验法[10],全局大气光和介质透射图都是通过手工制作的先验来评估的。这些先验可以区分雾霾区域和无雾霾区域。一些基于学习的方法 [29, 26, 35, 21, 11] 也遵循传统的去雾程序:估计透射图和大气光,并根据大气散射模型恢复无雾图像。然而,在很多非均匀天气条件下,灰霾和透射图都无法准确估算,这反过来又会影响后续的除杂,导致色彩过度和细节缺失等不良结果。现有方法在浓雾和非均匀雾度条件下面临挑战。为了克服物理模型带来的参数限制,需要一种端到端模型,它可以直接学习输入的雾霾图像到清晰的无雾霾图像映射,而不需要物理模型的支持。
由于很难同时获得雾霾图像和相应的地面实况图像对,真实场景中去雾霾训练集的数量非常有限,因此大多数研究都使用合成数据集作为训练集。然而,合成算法生成的合成数据集与真实雾霾场景的差异较大,而真实雾霾场景训练集的样本量又太小,因此我们需要更精确、更强大的先验来将额外的信息整合到深度除霾模型中。编码器-解码器架构是一个不错的选择。我们提出的三叉戟去噪网络(TDN)引入了编码器-解码器架构,作为重建粗略无雾霾图像的主网。此外,预训练编码器还能大大加快网络的收敛速度。
然而,由于大多数骨干网络(如 ResNet、DenseNet 等)中的池化层和降采样层无法学习,因此会丢失高频细节,而且重建图像中输入雾霾图像中雾霾密集的区域会非常模糊。为了解决这个问题,获得感知质量更好的无雾霾重建图像,我们引入了细节细化子网(DRN),以帮助主网获得更清晰、更忠实的结果。我们还制作了一个新颖的 FFT 损失函数来监督频域信息。有了 FFT 损失函数,对不同频段的监督更加统一,高频信息也更容易捕捉。此外,为了帮助适应非均质雾霾场景,防止在不同雾霾密度区域出现过度雾霾或雾霾不足的情况,我们添加了一个 U-Net 风格的雾霾密度图生成子网(HDMGN),以自动学习输入雾霾图像到雾霾密度图的映射。有了这两个子网,我们提出的三叉戟去雾网络不仅能适应密集雾霾场景,还能适应非均匀雾霾场景。大量的消融实验表明,所提出的 TDN 和 FFT 损失中的三个子网都很有效。在广泛使用的除霾基准 RESIDE 和真实世界的 NTIRE 除霾挑战测试集上,将 PSNR、SSIM 和 LPIPS 指标与之前的先进方法进行比较,实验表明 TDN 超越了所有之前的方法,具有更好的保真度和感知性。
总体而言,我们的贡献主要体现在以下三个方面:
- 我们提出了一种新颖的 FFT 损失函数来监督频域信号。结合空间域损耗和频域损耗,对不同频段的监督更加均匀,更容易捕捉高频信息。
- 我们提出了一种新颖的从粗到细的Trident去雾网络(TDN)模型,用于端到端学习从有雾到无雾的图像映射,并自动识别雾霾密度。我们提出的细节细化子网络可以细化具有更多高频细节的粗特征图,而我们提出的雾霾密度图生成子网络则可以自动重建雾霾密度图,无需额外监督即可实现非均匀去雾霾。
- 我们提出的TDN在NTIRE2020非均匀除雾挑战赛中获得第一名。在常用的去雾基准上的实验结果表明,TDN不仅在人工数据集上,而且在真实的雾天场景上都具有更好的逼真度和感知性,并且在高浓度雾天和非均匀雾天场景上都具有很好的泛化能力。
2. Trident去雾网络
在本节中,我们首先概述了Trident Dehazing Network(TDN),然后介绍了一些用于训练TDN的新型损失函数,这些损失函数可能对其他去雾网络有用。
2.1.网络架构
本文提出了一种三叉戟去雾网络(TDN),直接学习从输入的真实的世界非均匀模糊图像到无模糊清晰图像的映射。如图1所示,TDN由三个子网组成,编码器-解码器子网(EDN),细节细化子网(DRN)和雾度密度图生成子网(HDMGN),每个子网都用于特定目的:EDN重建无模糊图像的粗糙特征,DRN补充无模糊图像特征的高频细节,而HDMGN帮助获得输入模糊图像的不同区域中的模糊密度。可变形[38]卷积块从三个子网的级联特征图中获得最终的清晰输出。
图1:提出的Trident Dehazing Network(TDN)、细节细化子网络(DRN)和编码器-解码器子网络(EDN)的架构。分别表示张量加法和张量乘法。TDN由三个子网组成:EDN,DRN和HDMGN。然后将三个子网输出的雾度密度图和中间特征图连接起来,并输入到尾部可变形[38]卷积块以获得清晰的输出。
图2:所提出的雾度密度图生成子网络(HDMGN)的架构。“TConv”是“转置卷积”的缩写。
Encoder-Decoder sub-Net。在提出的EDN中,编码器部分支持ResNet风格的骨干,例如ResNet [18],ResNext [32],Res 2Net [16],DPN [13]等。它应该包含头层,支持4倍下采样和浅层特征提取,一层没有下采样,在第一个瓶颈中使用了三层下采样算子。我们使用在ImageNet 1 K中预训练的DPN 92作为编码器部分的骨干。它是一个功能强大的特征提取器,其中所提取的浅特征表示低级视觉特征,并且所提取的深特征能够捕获语义信息。EDN在训练期间固定编码器以利用预训练的DPN的能力作为“先验”,而编码器对于小样本训练集(如Dense Haze、NH-Haze等)带来了大量的额外信息。
解码器由5个可变形上采样块(DUB)组成,如图1(右下)所示。输入特征首先送入3×3可变形卷积块,然后与输出特征级联。级联后的特征送入1×1可变形卷积块和最近的上采样2×层,以获得上采样特征作为下一个DUB的输入特征。
我们分别通过连接(cat)特征图,从第二层和第三层的第一个下采样块的输出添加跳过连接到DUB 2,3的输入。我们使用可训练的实例归一化[30]来跳过连接。我们的编码器-解码器子网具有大容量,跳过连接使信息平滑流动,轻松训练大型网络。
Haze Density Map Generation sub-Net(HDMGN)。如图2所示,我们使用pix 2 pix [19]网络中提出的简单U-Net架构来实现Haze Density Map的生成。与pix 2 pix网络中的U-Net不同,我们增加了一个3×3的尾部卷积层来细化输出。由于大小划分要求,U-Net中只有6个下采样和上采样算子,输入大小应被64整除。如图3所示,可视化雾霾密度图中的区域越绿,雾霾越多。注意,我们没有对雾霾密度图进行任何额外的监督。U-Net可以自动实现。
细节细化子网络(DRN):受基于预上采样的单幅图像超分辨率网络的启发,本文提出了一个细节细化子网络,用于对下采样后的4×因子进行非线性特征映射,并使用逆像素混洗层将特征映射从空间到深度进行变换(下采样/去子像素),像素洗牌层[28]用于将特征图从深度更改为空间(上采样/子像素)。如图1所示,三个宽激活块(WAB)提供4×下采样因子的非线性特征映射。在WAB中,有两个3×3卷积层(后面是批量归一化层)和[33]中提出的宽激活层。WAB的通道扩展因子为4。受[31]的启发,我们使用残差缩放,即在将残差添加到主路径之前通过乘以0和1之间的常数来缩小残差,以防止训练不稳定。添加细节细化子网,训练过程更稳定,并且最终输出可以从稍微模糊的重建结果增强到具有更清晰细节的更清晰的结果。
2.2.损失函数
给定有雾图像和真实无雾图像,我们打算通过最小化损失函数来学习网络参数,损失函数由三个不同的分量组成,每个分量用于特定目的。各个损失项描述如下:
空间域损失:这是图像重建任务中常用的标准损失函数,它监督空间域中的重建无模糊结果:
其中,、TDN()分别表示由所提出的TDN重构的第i个地面实况样本和无模糊样本。
频域损耗:这是一种新的损失函数,它在频域中监督重建的无模糊结果。输出是频域中的幅度和相位,快速傅里叶变换(FFT)损失通过L1损失函数监督幅度和相位:
其中Ai和Pi分别表示第i个图像样本的幅度和相位,FFT(·)表示快速傅立叶变换,α作为两者之间的权衡参数(在我们的实验中,α = 1)。为了实现,我们使用PyTorch,其中FFT运算符可以通过torch.fft实现,并且可以在GPU中轻松计算。注意,FFT在训练过程中带来非常有限的时间成本,而所提出的频域损失大大提高了重建图像的视觉感知质量,而没有任何推理成本。
阈值限制损失:TDN的输入和输出范围是从0到1。直接通过TDN输出可能会导致阈值溢出。我们在TDN的尾部使用BReLU层[11]将输出数据限制在范围[0,1]内:
其中,分别地,B(·)是指BReLU算子,Oi是指第i个图像样本的最终输出。
地面实况中的黑色和白色像素区域更容易通过去雾网络重建以溢出阈值,它们可以通过BReLU层固定。但也有其他值会溢出。通过最小化阈值限制损失,可以将这些离群值校正到正常范围。
总损失函数定义为:
其中α、β和γ用作权衡参数以平衡不同的损耗项。
3.实验
在本节中,我们首先描述实验设置,即数据集,评估指标和实现细节,然后分别给出消融研究和模型分析。最后,我们进行实验来评估我们的性能以及竞争对手。
3.1.实验设置
数据集。在本文中,我们在两种类型的数据集上评估我们的网络:合成数据集和真实世界数据集。对于合成数据,我们选择了一个众所周知的代表性RESIDE [22]数据集。对于真实世界的数据,我们使用NTIRE 2018 Dehazing Challenge [1]中使用的I-HAZE数据集[7],O-HAZE数据集[6],NTIRE 2019 Dehazing Challenge [2]中使用的Dense Haze数据集[3]和NH-Haze数据集[8]进行实验,4]用于NTIRE 2020去雾挑战[5,9]。
RESIDE是一个广泛应用于去雾任务的标准数据集,它由室内和室外,合成和真实世界的雾图像组成。它总共有五个子集:室内训练集(ITS)、综合客观测试集(SOTS)、混合主观测试集(HSTS)、室外训练集(OTS)和真实的世界任务驱动测试集(RTTS)(真实图像)。大气光A是随机选择的,范围为(0.7,1.0),并且散射系数β的范围在(0.6,1.8).在我们的工作中,我们使用由13990幅合成图像组成的ITS来训练我们的网络,SOTS的室内部分被用作我们的测试集,其中包括500个室内图像。此外,我们在NTIRE 2018,NTIRE 2019和NTIRE 2020去雾挑战中介绍的真实世界去雾基准上实现了我们的方法,以进行进一步的评估。I-HAZE,O-HAZE,DenseHaze和NH-Haze包含25个室内烟雾图像,35个室外烟雾图像,45个密集模糊图像和45个非均匀模糊图像及其相应的地面实况(GT)分别用于训练,5个hazy-gt对用于验证,5个用于测试。NTIRE挑战数据集使用专业的hazy/模拟雾霾场景真实的条件的雾发生器。详细信息如表1所示。对于I-HAZE、O-HAZE和Dense-Haze数据集,我们使用训练集进行训练,使用验证集进行测试。对于NH-Haze数据集,因为验证集现在还没有公开,所以我们使用图像1 - 40作为消融研究的训练集,使用图像41 - 45作为消融研究的测试集。
表1:本文中使用的去雾基准的细节。“I”和“O”分别表示“室内”和“室外”。“*”表示地面实况图像(GT)不是公开的。I-HAZE和O-HAZE的图像大小是平均估计值。“r”表示模糊图像是真实的世界图像还是合成图像。
对于定性评估,我们根据两个评估指标来衡量我们的方法的结果:峰值信噪比(PSNR)和结构相似度指数(SSIM)是评价图像质量的重要指标,因为NTIRE 2020去雾挑战的目的是产生具有最佳感知质量并且类似于参考地面真实的高质量结果,我们为NH-雾度数据集添加了感知测量学习感知图像块相似性(LPIPS)[36],并将我们的去雾效果与主观视觉效果进行了比较。
模型设置和实现细节。在训练过程中,从训练图像中裁剪大小为256 × 256的补丁。我们在RGB通道中训练TDN,并通过随机旋转90,180,270度和水平翻转来增强训练数据集。我们的模型使用Adam优化器进行优化,β1 = 0.9,β2 = 0.999,并且,为了使损失函数最小化,将初始学习率设置为10−4,然后在第30、55和80个epoch中降低到一半。批量大小(batchsize)为40,总epoch数为100,每个epoch有400次迭代。我们使用PyTorch来实现我们的模型,其中两个GTX 1080 Ti用于训练,一个用于测试。训练时间约为一天,TDN中使用的骨干网络是预训练的DPN 92,它首先在ImageNet 5 K中进行预训练,然后在ImageNet 1 K中进行微调。TDN的其他参数使用随机权重进行初始化。DPN 92的头层,层1和层2的学习速度是整个网络的一半,DPN 92的layer 3、layer 4-block 1和可变形卷积偏移掩模层的学习速率分别是整个网络的0.7、0.9和0.3倍。在宽激活块中使用的残差标度为0.2。我们根据经验将损失加权因子α、β和γ设置为0.5、0.5和1。
3.2.消融研究和模型分析
三个子网的有效性。由于编码器-解码器子网(EDN)的编码器部分的池化层,EDN更倾向于学习粗糙的低频信息,得到有点模糊的结果。细节细化子网(DRN)被提出来通过下采样4×因子的非线性特征映射来细化高频细节。为了适应非均匀图像,我们添加了一个U-Net风格的架构来学习模糊输入图像到模糊密度图映射。简而言之,Trident Dehazing Network是一个从粗到细的架构,具有自动模糊密度识别功能,用于非均匀去雾。
我们提供了三个子网的输出的可视化图来证明三个子网的有效性。如图3所示,EDN的输出特征图重建了无模糊图像的粗糙特征,其中包含图像的大致轮廓信息。请注意,一些极端离群点像素会导致EDN可视化图中存在一些伪影。它将由尾部可变形卷积层修复。尾部可变形卷积层可以自动忽略这些离群点,并通过HDMGN精确地重建了雾密度的分布。在DRN的输出特征图中,边缘信息和高频细节是主要成分。三个子网的消融研究结果如表2(5)-(8)所示。消融研究中使用了更快速推断的主干DPN 68。设置(5)是EDN和整个网络的最终设置,模型(5)中只有主网EDN,模型(6)和模型(7)分别丢失了HDMGN和DRN,模型(8)包含所有子网,如果删除任何子网,性能会更差。
图3:从左到右分别是编码解码子网络(EDN)、雾度密度图生成子网络(HDMGN)和细节细化子网络(DRN)的可视化图。HDMGN的输出是3通道特征图,经过[0,1]钳制后可视化为RGB图。为了可视化EDN和DRN的输出,我们将EDN和DRN的输出特征图在通道维度上求和,并使用sigmoid层来获得可视化图。输入模糊图像“52.png”可以在图1中找到.
可变形卷积层的有效性。如图4左图所示,在没有可变形卷积的情况下,局部区域被识别为相同级别的雾度密度,这导致当一些对象成为前景区域,而被雾度遮挡的其他部分成为背景区域时,雾度部分无法重建,而前景部分被过度锐化。当使用可变形卷积时,由于2D偏移极大地增强了CNN变换建模的能力,因此前景和模糊背景区域特征被很好地分离以提供不同程度的去雾。
图4:左:示例输出图像w/o可变形卷积。中间:具有可变形卷积的示例输出图像。右侧:包含伪影的示例输出图像w/o可变形卷积。
可变形卷积层不仅可以帮助识别雾霾的密度,还可以修复异常像素造成的伪影。如图4右和图3左所示,虽然我们做了[0,1]箝位来获得输出,但也有一些极端的异常像素会导致中间特征图和最终输出中的伪影。通过将解码器中的标准卷积层和尾卷积块替换为可变形卷积层,可变形卷积层中的可学习卷积偏移忽略这些离群值并选择适合的卷积块。如表2设置(4)和设置(5)中所示的消融研究结果也证明了可变形卷积层的有效性。
表2:损失函数和架构的消融研究结果。“DC”表示TDN将可变形卷积用于EDN的解码器部分和尾卷积层而不是标准卷积层。“IN”表示EDN将实例归一化层添加到跳过连接。L1,FFT和BReLU指的是空间域损失,频域损失和阈值限制损失,结果最好的行用粗体表示。
实例归一化层的消融研究和提出的损失函数。我们提供了一个示例图像的定性视觉效果比较,如图5所示。在跳过连接中添加实例归一化层后,树叶和草地的颜色饱和度增加,从而获得更好的可视化重建无模糊图像。表2设置(3)和设置(4)的比较证明了具有实例归一化层的网络的结果获得了更好的感知质量,并且与参考地面真实值更相似。
图5:左:模糊图像“48.png”。中:从TDN中删除实例规范化层的网络的模糊图像。右:TDN的模糊图像。
如图6所示,使用所有提出的损失函数的TDN获得了最佳的损失收敛结果。阈值限制(BReLU)损失有助于在早期和后期训练网络。添加频域(FFT)损失阻碍了空间域的训练过程,在早期,但当整个训练过程结束时,空间域(L1)损失曲线收敛到一个更好的结果,与L1曲线和L1 BReLU曲线相比。比较表2设置(1)到(3),我们可以发现,加入阈值限制损失后,性能稍好。此外,增加频域损失,无论是对地面真实值的恢复保真度(PSNR和SSIM)还是感知质量(LPIPS),重建的无模糊图像都获得了更好的性能。
图6:L1 loss-Epoch消融研究曲线,设置(1)~(3)L1:仅使用空间域损失作为损失函数L1 BReLU:使用空间域损失和阈值限制损失作为损失函数L1 BReLU FFT:使用空间域损失、阈值限制损失和频域损失作为损失函数
3.3.与最先进方法的比较
本节展示了我们提出的TDN与最先进的方法在真实世界基准数据集I-HAZE,O-HAZE,Dense-Haze和NH-Haze以及合成数据集ITS上的比较。我们在相应的训练集上重新训练了不同基准的最先进的模型,以进行公平的比较。报告了定量评估和视觉效果。
对比方法。比较中包括六种最先进的方法:DCP [17],DehazeNet [11],AOD-Net [21],DCPDN [34],GCANet [12]和FFA [25]。
定量结果。定量比较结果如表3所示。在综合测试集RESIDE(ITS)上,我们的方法仅次于FFA。注意FFA有太多的块,并且在整个图像的尺度上进行所有特征映射运算,这导致了巨大的计算量,而我们的TDN在下采样尺度上进行主要卷积运算,该方法具有快速的推理时间,一幅1600×1200图像的推理时间仅为0.64s,在真实的NTIRE去雾挑战数据集上,该方法在PSNR和SSIM方面都具有最好的性能,特别是在高雾度和非均匀雾度场景中,该方法的性能明显优于其他方法。
视觉效果对比。如图7至图11所示,DCP在ITS测试集上得到的结果较暗,在真实的NTIRE去雾挑战测试集上得到的结果较蓝。在NTIRE测试集上,DehazeNet中丢失了大部分颜色信息,同时产生了一些伪影。AOD-Net无法有效去除雾。在O-Haze测试集上,DCPDN恢复图像的颜色有很大的偏差。在密集的雾测试和NH-Haze数据集,DCPDN产生非常不愉快的文物,FFA得到肮脏的结果,而GCANet过度去雾模糊的图像,导致非常黑暗的错误结果,结果表明,DCPDN和GCANet算法对高浓度或非均匀的雾霾场景鲁棒性较差,FFA算法在I-HAZE和O-HAZE测试集上取得了较好的效果,但由于占用了大量的图形内存,只能使用chop和concatenation策略进行推理,导致严重的棋盘状伪影,只有我们提出的TDN在所有常用的去雾基准测试中重建了忠实和清晰的无模糊结果,伪影很小,感知质量很好。
NTIRE-2020 Dehazing Challenge。对于新发布的NTIRE 2020-Dehaze数据集,图像中呈现的雾度比文献中的正常图像更加不均匀。如图10所示,最先进的方法由于雾度不均匀地覆盖整个图像,且不同区域的雾度不同,因此性能严重下降。通过从非均匀雾度图像中提取密度信息,由TDN生成的去雾图像在视觉上更令人愉悦。我们在NTIRE 2020数据集41 - 45上评估了这些方法的定量性能,因为验证集的地面实况图像现在不可用。如表3所示(NH-Haze),TDN优于所有其他最先进的方法。
表4包括竞赛中排名前6的感知质量方法。我们发现,TDN是NTIRE 2020-Dehazing Challenge中排名前6的感知质量方法之一。
4.结论
本文提出了一种新的由粗到精的Trident Dehazing Network(TDN)模型,TDN由三个子网络组成,编码器-解码器子网络重构粗糙的无雾特征,细节细化子网络对编码器中池化层丢失的高频细节进行细化,烟雾密度图生成子网络可以自动重建烟雾密度图,而无需额外的监督。大量的实验结果表明,TDN是强大的,不仅对合成数据集,但也对现实世界的场景与密集的雾和非均匀的雾,并优于最先进的更好的保真度和感知。
表3:不同方法在SOTS室内、I-HAZE、O-HAZE、Dense Haze和NH-Haze测试集上的PSNR/SSIM/LPIPS。
表4:NTIRE 2020测试数据集上顶级感知质量方法的平均PSNR/SSIM/LPIPS/MOS排名。
具有最佳结果的行以粗体显示。
【这篇文章的创新点就是设计一个网络TDN:一张有雾图经过三个部分EDN、DRN、HDMGN三个模块的输出与中间特征图concat起来,并输入到尾部可变形卷积块以获得清晰的输出;并对损失函数进行了改进。】