A Transformer-Based Diffusion Model for All-in-One Weather-Degraded Image Restoration

doi:10.16183/j.cnki.jsjtu.2023.043

Abstract

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

CLC Number:

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.

Figures/Tables 15

Fig.1

Fig.2

Fig.3

Fig.4

Tab.1

Tab.2

Quantitative comparison of different methods in all-in-one weather-degraded image restoration tasks

方法	源	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

Tab.2

Fig.5

Fig.6

Fig.7

Tab.3

Quantitative comparison of different methods on Snow100K dataset

方法	源	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

Tab.3

Tab.4

Quantitative comparison of different methods on Test1 dataset

方法	源	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

Tab.4

Tab.5

Quantitative comparison of different methods on Raindrop-A dataset

方法	源	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

Tab.5

Tab.6

Quantitative comparison of different methods on natural weather-degraded image dataset

方法	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

Tab.6

Fig.8

Tab.7

Ablation experimental results of individual components of noise estimation network

网络模型	网络组成								单步估计用时/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

Tab.7

References 34

[1]	高涛, 文渊博, 陈婷, 等. 基于窗口自注意力网络的单图像去雨算法[J]. 上海交通大学学报, 2023, 57(5): 613-623. doi: 10.16183/j.cnki.jsjtu.2022.032 URL
	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⁶