面向多天气退化图像恢复的自注意力扩散模型

doi:10.16183/j.cnki.jsjtu.2023.043

摘要/Abstract

摘要：

复杂天气下的图像恢复对后续高级计算机视觉任务具有重要意义.然而,多数现有图像恢复算法仅能去除单一天气退化,鲜有针对多天气退化图像恢复的同一模型.对此,结合去噪扩散概率模型和视觉Transformer,提出一种用于多天气退化图像恢复的自注意力扩散模型.首先,利用天气退化图像作为条件来引导扩散模型反向采样生成去除退化的干净背景图像.其次,提出次空间转置自注意力噪声估计网络,利用退化图像和噪化状态来估计噪声分布,包括次空间转置自注意力机制 (STSA) 和双分组门控前馈网络 (DGGFFN).STSA利用次空间变换系数实现有效学习特征全局性长距离依赖的同时,可显著降低计算负担;DGGFFN利用双分组门控机制来增强前馈网络的非线性表征能力.实验结果表明,在5个天气退化图像数据集上,相比近来同类算法All-in-One和TransWeather,本文算法所得恢复图像的平均峰值信噪比分别提高3.68和3.08 dB,平均结构相似性分别提高2.93%和3.13%.

关键词: 计算机视觉, 扩散模型, 图像恢复, Transformer, 天气退化图像

Abstract:

Image restoration under adverse weather conditions is of great significance for the subsequent advanced computer vision tasks. However, most existing image restoration algorithms only remove single weather degradation, and few studies has been conducted on all-in-one weather-degraded image restoration. The denoising diffusion probability model is combined with Vision Transformer to propose a Transformer-based diffusion model for all-in-one weather-degraded image restoration. First, the weather-degraded image is utilized as the condition to guide the reverse sampling of diffusion model and generate corresponding clean background image. Then, the subspace transposed Transformer for noise estimation (NE-STT) is proposed, which utilizes the degraded image and the noisy state to estimate noise distribution, including the subspace transposed self-attention (STSA) mechanism and a dual grouped gated feed-forward network (DGGFFN). The STSA adopts subspace transformation coefficient to effectively capture global long-range dependencies while significantly reducing computational burden. The DGGFFN employs the dual grouped gated mechanism to enhance the nonlinear characterization ability of feed-forward network. The experimental results show that in comparison with the recently developed algorithms, such as All-in-One and TransWeather, the method proposed obtains a performance gain of 3.68 and 3.08 dB in average peak signal-to-noise ratio while 2.93% and 3.13% in average structural similarity on 5 weather-degraded datasets.

Key words: computer vision, diffusion model, image restoration, Transformer, weather-degraded image

中图分类号:

秦菁, 文渊博, 高涛, 刘瑶. 面向多天气退化图像恢复的自注意力扩散模型[J]. 上海交通大学学报, 2024, 58(10): 1606-1617.

QIN Jing, WEN Yuanbo, GAO Tao, LIU Yao. A Transformer-Based Diffusion Model for All-in-One Weather-Degraded Image Restoration[J]. Journal of Shanghai Jiao Tong University, 2024, 58(10): 1606-1617.

图/表 15

图1

图2

图3

图4

表1

表2

不同算法在多天气退化图像恢复任务上的定量对比

方法	源	Snow100K-L PSNR/SSIM	Snow100K-M PSNR/SSIM	Snow100K-S PSNR/SSIM	Test1 PSNR/SSIM	Raindrop-A PSNR/SSIM
Uformer^[10]	CVPR 2022	26.24/0.8680	32.11/0.9316	34.00/0.9445	16.32/0.7565	30.33/0.9335
Restormer^[11]	CVPR 2022	$29.57 =$ / $0.9106 =$	$33.71 =$ / $0.9490 =$	$35.43 =$ / $0.9583 =$	$29.81 =$ / $0.9208 =$	$31.10 =$ / $0.9337 =$
All-in-One^[3]	CVPR 2020	28.14/0.8901	30.96/0.9290	32.63/0.9392	25.87/ $0.8996_$	$31.35_$ / $0.9299_$
TransWeather^[4]	CVPR 2022	$29.15_$ / $0.8930_$	$32.02_$ / $0.9343_$	$32.98_$ / $0.9447_$	$28.61_$ /0.8893	30.21/0.9179
AWIR-TDM		31.69/0.9240	35.47/0.9565	37.16/0.9642	31.68/0.9347	32.33/0.9429

表2

图5

图6

图7

表3

不同算法在雪数据集Snow100K上的定量对比

方法	源	Snow100K-L PSNR/SSIM	Snow100K-M PSNR/SSIM	Snow100K-S PSNR/SSIM	平均指标 PSNR/SSIM
RESCAN^[22]	ECCV 2018	26.08/0.8108	29.95/0.8860	31.51/0.9032	29.28/0.8667
SPANet^[23]	CVPR 2019	23.70/0.7930	28.06/0.8680	29.92/0.8260	27.23/0.8290
DesnowNet^[14]	TIP 2018	27.17/ $0.8983_$	30.87/0.9409	32.33/0.9500	30.12/ $0.9300_$
JSTASR^[24]	ECCV 2020	25.32/0.8076	29.11/0.8843	31.40/0.9012	28.61/0.8644
SwinIR^[25]	CVPR 2021	28.18/0.8800	31.42/0.9284	33.96/ $0.9567_$	31.19/0.9217
DDMSNet^[26]	TIP 2021	28.85/0.8772	32.89/0.9330	34.34/0.9445	32.03/0.9182
TransWeather^[4]	CVPR 2022	$29.21_$ /0.8947	$33.41_$ / $0.9416_$	$34.92_$ /0.9543	$32.51_$ /0.9279
AWIR-TDM		30.64/0.9193	35.26/0.9472	37.05/0.9680	34.32/0.9448

表3

表4

不同算法在雨雾数据集Test1上的定量对比

方法	源	Test1
方法	源	PSNR/SSIM
CycleGAN^[27]	ICCV 2017	17.62/0.6560
pix2pix^[28]	CVPR 2017	19.09/0.7100
HRGAN^[20]	CVPR 2019	21.56/0.8550
SwinIR^[25]	CVPR 2021	23.23/0.8685
PCNet^[29]	TIP 2021	26.19/0.9015
MPRNet^[30]	CVPR 2021	28.03/0.9192
Restormer^[11]	CVPR 2022	$28.44_$ / $0.9263_$
AWIR-TDM		29.13/0.9428

表4

表5

不同算法在雨滴数据集Raindrop-A上的定量对比

方法	源	Raindrop-A
方法	源	PSNR/SSIM
pix2pix^[28]	CVPR 2017	28.02/0.8547
Attn. GAN^[17]	CVPR 2018	31.59/0.9170
DuRN^[31]	CVPR 2019	31.24/0.9259
RaindropAttn^[32]	ICCV 2019	31.44/0.9263
SwinIR^[25]	CVPR 2021	30.82/0.9035
CCN^[33]	CVPR 2021	31.34/ $0.9500_$
IDT^[34]	TPAMI 2022	$31.87_$ /0.9313
AWIR-TDM		32.84/0.9571

表5

表6

不同算法在自然天气退化图像数据集上的定量对比

方法	Snow NIQE/SSEQ/NIMA	RainMist NIQE/SSEQ/NIMA	RainStreak NIQE/SSEQ/NIMA	Raindrop NIQE/SSEQ/NIMA
Uformer^[10]	3.395/28.31/2.644	4.021/26.88/3.289	3.771/27.19/3.376	4.792/34.06/4.260
Restormer^[11]	3.267/27.90/2.570	3.912/ $24.30_$ / $3.874_$	$3.694_$ /26.47/ $3.480_$	4.658/ $30.72_$ / $4.317_$
All-in-One^[3]	3.561/29.62/2.384	4.253/25.20/3.419	3.895/27.03/3.352	$5.000_$ /38.94/4.073
TransWeather^[4]	3.020/ $27.78_$ / $2.793_$	$3.791_$ /24.62/3.700	3.765/27.11/3.282	4.702/31.81/4.213
AWIR-TDM	$3.134_$ /27.62/2.980	3.752/24.23/3.965	3.636/ $26.75_$ /3.491	4.647/30.46/4.328

表6

图8

表7

噪声估计网络的组件消融实验结果

网络模型	网络组成								单步估计用时/s	Raindrop-A
网络模型	ResBlock	ViT	SA	TSA	STSA	FFN	SGFFN	DGGFFN	单步估计用时/s	PSNR / SSIM
基线网络	√		√						0.6634	29.17/0.9162
ⅰ		√	√			√			1.0725	29.86/0.9203
ⅱ		√		√		√			0.8603	31.54/0.9297
ⅲ		√			√	√			0.3961	30.49/0.9381
ⅳ		√			√		√		0.4033	$31.73_$ / $0.9394_$
AWIR-TDM		√			√			√	$0.4080_$	32.33/0.9429

表7

参考文献 34

[1]	高涛, 文渊博, 陈婷, 等. 基于窗口自注意力网络的单图像去雨算法[J]. 上海交通大学学报, 2023, 57(5): 613-623. doi: 10.16183/j.cnki.jsjtu.2022.032
	GAO Tao, WEN Yuanbo, CHEN Ting, et al. A single image deraining algorithm based on Swin Transformer[J]. Journal of Shanghai Jiao Tong University, 2023, 57(5): 613-623.
[2]	黄鹤, 胡凯益, 李战一, 等. 融合MCAP和GRTV正则化的无人机航拍建筑物图像去雾方法[J]. 上海交通大学学报, 2023, 57(3): 613-623.
	HUANG He, HU Kaiyi, LI Zhanyi, et al. An image dehazing method for UAV aerial photography to buildings combining MCAP and GRTV regularization[J]. Journal of Shanghai Jiao Tong University, 2023, 57(3): 613-623.
[3]	LI R, ROBBY T T, LOONG-FAH C. All in one bad weather removal using architectural search[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Seattle, WA, USA: IEEE, 2020: 3175-3185.
[4]	VALANARASU J M J, YASARLA R, PATEL V M. Transweather: Transformer-based restoration of images degraded by adverse weather conditions[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. NewOrleans, LA, USA: IEEE, 2022: 2353-2363.
[5]	GOODFELLOW I, POUGET-ABADIE J, MIRZA M, et al. Generative adversarial networks[J]. Communications of the ACM, 2020, 63(11): 139-144.
[6]	KINGMA D P, WELLING M. Auto-encoding variational bayes[DB/OL]. (2013-12-20)[2023-02-06]. https://arxiv.org/abs/1312.6114.
[7]	HO J, JAIN A, ABBEEL P. Denoising diffusion probabilistic models[J]. Advances in Neural Information Processing Systems, 2020, 33: 6840-6851.
[8]	DHARIWAL P, NICHOL A. Diffusion models beat gans on image synthesis[J]. Advances in Neural Information Processing Systems, 2021, 34: 8780-8794.
[9]	PEEBLES W, XIE S. Scalable diffusion models with Transformers[DB/OL]. (2022-12-19)[2023-02-06]. https://arxiv.org/abs/2212.09748.
[10]	WANG Z, CUN X, BAO J, et al. Uformer: A general u-shaped transformer for image restoration[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. NewOrleans, LA, USA: IEEE, 2022: 17683-17693.
[11]	ZAMIR S W, ARORA A, KHAN S, et al. Restormer: Efficient transformer for high-resolution image restoration[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. NewOrleans, LA, USA: IEEE, 2022: 5728-5739.
[12]	YAO T, LI Y, PAN Y, et al. Dual vision transformer[DB/OL]. (2022-07-11) [2023-02-06]. https://arxiv.org/abs/2207.04976.
[13]	CHEN L, CHU X, ZHANG X, et al. Simple baselines for image restoration[C]// Proceedings of the European Conference on Computer Vision. Tel Aviv, Israel: Springer, 2022: 17-33.
[14]	LIU Y F, JAW D W, HUANG S C, et al. DesnowNet: Context-aware deep network for snow removal[J]. IEEE Transactions on Image Processing, 2018, 27(6): 3064-3073.
[15]	鲍先富, 强赞霞, 杨关. 功能解耦和谱特征融合的雪霾消除模型[J]. 计算机工程与应用, 2023, 59(13): 211-219. doi: 10.3778/j.issn.1002-8331.2203-0566
	BAO Xianfu, QIANG Zanxia, YANG Guan. Generative adverbial network for function decoupling and edge feature fusion for snow and haze elimination[J]. Computer Engineering & Applications, 2023, 59(13):211-219.
[16]	柴国强, 王大为, 芦宾, 等. 基于注意机制的轻量化稠密连接网络单幅图像去雨[J]. 北京航空航天大学学报, 2022, 48(11): 2186-2192.
	CHAI Guoqiang, WANG Dawei, LU Bin, et al. Lightweight densely connected network based on attention mechanism for single-image deraining[J]. Journal of Beijing University of Aeronautics & Astronautics, 2022, 48(11): 2186-2192.
[17]	QIAN R, TAN R T, YANG W, et al. Attentive generative adversarial network for raindrop removal from a single image[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City, UT, USA: IEEE, 2018: 2482-2491.
[18]	CHEN H, WANG Y, GUO T, et al. Pre-trained image processing transformer[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Kuala Lumpur, Malaysia: IEEE, 2021: 12299-12310.
[19]	LI B, LIU X, HU P, et al. All-in-one image restoration for unknown corruption[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. NewOrleans, LA, USA: IEEE, 2022: 17452-17462.
[20]	LI R, CHEONG L F, TAN R T. Heavy rain image restoration: Integrating physics model and conditional adversarial learning[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Long Beach, CA, USA: IEEE, 2019: 1633-1642.
[21]	LI S, ARAUJO I B, REN W, et al. Single image deraining: A comprehensive benchmark analysis[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Long Beach, CA, USA: IEEE, 2019: 3838-3847.
[22]	LI X, WU J, LIN Z, et al. Recurrent squeeze-and-excitation context aggregation net for single image deraining[C]// Proceedings of the European Conference on Computer Vision. Salty Lake City, UT, USA: Springer, 2018: 254-269.
[23]	WANG T, YANG X, XU K, et al. Spatial attentive single-image deraining with a high quality real rain dataset[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Long Beach, CA, USA: IEEE, 2019: 12270-12279.
[24]	CHEN W, FANG H, DING J, et al. JSTASR: Joint size and transparency-aware snow removal algorithm based on modified partial convolution and veiling effect removal[C]// Proceedings of the European Conference on Computer Vision. Glasgow, UK: Springer, 2020: 754-770.
[25]	LIANG J, CAO J, SUN G, et al. Swinir: Image restoration using swin transformer[C]// Proceedings of the IEEE/CVF International Conference on Computer Vision. Montreal, canada: IEEE, 2021: 1833-1844.
[26]	ZHANG K, LI R, YU Y, et al. Deep dense multi-scale network for snow removal using semantic and depth priors[J]. IEEE Transactions on Image Processing, 2021, 30: 7419-7431.
[27]	ZHU J Y, PARK T, ISOLA P, et al. Unpaired image-to-image translation using cycle-consistent adversarial networks[C]// Proceedings of the IEEE International Conference on Computer Vision. Venice, Italy: IEEE, 2017: 2223-2232.
[28]	ISOLA P, ZHU J Y, ZHOU T, et al. Image-to-image translation with conditional adversarial networks[C]// Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. Venice, Italy: IEEE, 2017: 1125-1134.
[29]	JIANG K, WANG Z, YI P, et al. Rain-free and residue hand-in-hand: A progressive coupled network for real-time image deraining[J]. IEEE Transactions on Image Processing, 2021, 30: 7404-7418.
[30]	ZAMIR S W, ARORA A, KHAN S, et al. Multi-stage progressive image restoration[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Kuala Lumpur, Malaysia: IEEE, 2021: 14821-14831.
[31]	LIU X, SUGANUMA M, SUN Z, et al. Dual residual networks leveraging the potential of paired operations for image restoration[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Long Beach, CA, USA: IEEE, 2019: 7007-7016.
[32]	QUAN Y, DENG S, CHEN Y, et al. Deep learning for seeing through window with raindrops[C]// Proceedings of the IEEE/CVF International Conference on Computer Vision. Seoul Korea: IEEE, 2019: 2463-2471.
[33]	QUAN R, YU X, LIANG Y, et al. Removing raindrops and rain streaks in one go[C]// Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. Kuala Lumpur, Malaysia: IEEE, 2021: 9147-9156.
[34]	XIAO J, FU X, LIU A, et al. Image de-raining transformer[J]. IEEE Transactions on Pattern Analysis & Machine Intelligence, 2022: 1-18.

自注意力方法	乘法计算量	加法计算量
SA	1.07×10⁹	1.06×10⁹
TSA	8.39×10⁶	8.26×10⁶
STSA	2.62×10⁶	2.49×10⁶