[论文解读] DW-GAN: A Discrete Wavelet Transform GAN for NonHomogeneous Dehazing

Minghan Fu, Liu Huan+3 more|arXiv (Cornell University)|Apr 18, 2021

Image Enhancement Techniques参考文献 41被引用 4

一句话总结

DW-GAN 提出了一种用于非均匀图像去雾的双分支生成对抗网络，通过集成二维离散小波变换（DWT）以保留高频细节，并利用 ImageNet 预训练的 Res2Net 实现知识蒸馏，在训练数据有限的情况下显著提升了性能。该方法在合成数据集和真实世界数据集上均取得了最先进结果，包括在 NTIRE2021 去雾挑战赛中取得 PSNR 21.08 dB 和 SSIM 0.8393 的最佳表现。

ABSTRACT

Hazy images are often subject to color distortion, blurring, and other visible quality degradation. Some existing CNN-based methods have great performance on removing homogeneous haze, but they are not robust in non-homogeneous case. The reasons are mainly in two folds. Firstly, due to the complicated haze distribution, texture details are easy to be lost during the dehazing process. Secondly, since the training pairs are hard to be collected, training on limited data can easily lead to over-fitting problem. To tackle these two issues, we introduce a novel dehazing network using 2D discrete wavelet transform, namely DW-GAN. Specifically, we propose a two-branch network to deal with the aforementioned problems. By utilizing wavelet transform in DWT branch, our proposed method can retain more high-frequency knowledge in feature maps. In order to prevent over-fitting, ImageNet pre-trained Res2Net is adopted in the knowledge adaptation branch. Owing to the robust feature representations of ImageNet pre-training, the generalization ability of our network is improved dramatically. Finally, a patch-based discriminator is used to reduce artifacts of the restored images. Extensive experimental results demonstrate that the proposed method outperforms the state-of-the-arts quantitatively and qualitatively.

研究动机与目标

• 为解决非均匀去雾的挑战，其中雾霾分布复杂且纹理细节易丢失。
• 通过利用 ImageNet 的预训练特征，在数据量有限的情况下提升模型泛化能力并减少过拟合。
• 通过在网络架构中引入离散小波变换，保留去雾过程中的高频分量（如边缘、纹理）。
• 通过基于图像块的对抗判别器生成逼真且无伪影的去雾图像。
• 在有限的训练样本下，实现在多样化真实世界和合成雾图数据集上的鲁棒性能。

提出的方法

• 提出一种双分支生成器架构：其中一支路采用基于 DWT 的下采样以保留频域信息并减少参数量。
• DWT 分支将标准卷积层替换为可学习的小波下采样模块，以在特征图中保留高频特征。
• 知识适配分支采用 ImageNet 预训练的 Res2Net 作为特征编码器，将先验知识注入模型，从而在小样本数据集上提升泛化能力。
• 集成通道注意力与空间注意力机制，帮助网络聚焦于关键特征和有雾区域。
• 采用基于图像块的判别器，以增强局部真实性并减少去雾输出中的伪影。
• 在数据混合过程中，对 NH-HAZE 和 NH-HAZE2 数据集应用伽马校正，以对齐亮度分布，提升训练稳定性。

实验结果

研究问题

• RQ1 将离散小波变换集成到基于 GAN 的去雾网络中，是否能有效提升非均匀雾霾图像中高频细节的保留？
• RQ2 通过 ImageNet 预训练的 Res2Net 实现知识蒸馏，在训练数据有限时，其提升模型泛化能力的效果如何？
• RQ3 所提出的双分支架构在真实世界和合成非均匀去雾基准测试中，是否在定量指标和感知质量方面均优于现有 SOTA 方法？
• RQ4 对训练数据进行伽马校正（γ=0.65）在合并具有不同亮度分布的数据集时，能在多大程度上提升模型性能？
• RQ5 所提出方法是否能在具有挑战性的现实世界数据集（如 DENSE-HAZE 和 NH-HAZE2）上实现鲁棒的去雾性能，并产生极少伪影？

主要发现

• DW-GAN 在 NTIRE2021 去雾挑战赛中取得 PSNR 21.08 dB 和 SSIM 0.8393 的成绩，位列顶尖方法之列。
• 在 NH-HAZE 和 NH-HAZE2 数据集上，DW-GAN 在定量指标和视觉质量方面均优于 FFA、GCANet 和 TDN 等 SOTA 方法。
• 在 DENSE-HAZE 数据集上，DW-GAN 有效去除了密集雾霾，颜色偏差极小，并成功保留了精细细节，优于 FFA 和 TDN。
• 对 NH-HAZE 数据集应用伽马校正（γ=0.65）后，其亮度分布与 NH-HAZE2 更为一致，降低了分布差异，从而提升了训练性能。
• DW-GAN 的推理时间快于 GCANet、FFA 和 TDN，同时保持了更优的性能，表明其在效率与准确率之间具有有利的权衡。
• 消融实验证实，DWT 分支和知识适配分支均对性能有显著贡献，注意力机制进一步提升了细节恢复能力。