Two‐stage single image dehazing network using swin‐transformer

IET Image ProcessingVolume 16, Issue 9 p. 2518-2534 ORIGINAL RESEARCHOpen Access Two-stage single image dehazing network using swin-transformer Xiaoling Li, Xiaoling Li School of Information and Electronic Engineering, Shandong Technology and Business University, Yantai, China Institute of Network Technology, ICT, Yantai, ChinaSearch for more papers by this authorZhen Hua, Corresponding Author Zhen Hua [email protected] orcid.org/0000-0003-1638-2974 School of Information and Electronic Engineering, Shandong Technology and Business University, Yantai, China Institute of Network Technology, ICT, Yantai, China Correspondence Zhen Hua, School of Information and electronic engineering, Shandong Technology and Business University, Yantai 264005, China. Email: [email protected]Search for more papers by this authorJinjiang Li, Jinjiang Li orcid.org/0000-0002-2080-8678 School of Computer Science and Technology, Shandong Technology and Business University, Yantai, ChinaSearch for more papers by this author Xiaoling Li, Xiaoling Li School of Information and Electronic Engineering, Shandong Technology and Business University, Yantai, China Institute of Network Technology, ICT, Yantai, ChinaSearch for more papers by this authorZhen Hua, Corresponding Author Zhen Hua [email protected] orcid.org/0000-0003-1638-2974 School of Information and Electronic Engineering, Shandong Technology and Business University, Yantai, China Institute of Network Technology, ICT, Yantai, China Correspondence Zhen Hua, School of Information and electronic engineering, Shandong Technology and Business University, Yantai 264005, China. Email: [email protected]Search for more papers by this authorJinjiang Li, Jinjiang Li orcid.org/0000-0002-2080-8678 School of Computer Science and Technology, Shandong Technology and Business University, Yantai, ChinaSearch for more papers by this author First published: 14 April 2022 https://doi.org/10.1049/ipr2.12506AboutSectionsPDF ToolsRequest permissionExport citationAdd to favoritesTrack citation ShareShare Give accessShare full text accessShare full-text accessPlease review our Terms and Conditions of Use and check box below to share full-text version of article.I have read and accept the Wiley Online Library Terms and Conditions of UseShareable LinkUse the link below to share a full-text version of this article with your friends and colleagues. Learn more.Copy URL Share a linkShare onFacebookTwitterLinkedInRedditWechat Abstract Hazy images often have color distortion, blur and other visible visual quality degradation, affecting the performance of some advanced visual tasks. Therefore, single image dehazing has always been a challenging and significant problem. Convolutional neural network has been widely used in image dehazing task, but the limitations of convolutional operation limit the development of dehazing task. Nowadays, Transformer offers a holistic approach to CV development and does not grow in location as the network deepens. For this reason, a hierarchical Transformer is introduced for use in the dehazing network. Specifically, the codec is improved and Transformer and CNN are combined to achieve basic feature extraction in the first stage. The encoder only models the global relationship at each layer, reducing the resolution of the feature map continuously and expanding the field of perception. In addition, an inter-block supervision mechanism is added between encoder unit and decoder unit to refine features and supervise and select them, thus improving the efficiency of feature transmission. In the second stage, the original resolution block is used to extract the local features, and then feature fusion and interaction are carried out. In addition, to ensure the authenticity of the transmission of characteristic signals in the first stage and improve the transmission efficiency of the network, fusion attention mechanism is added between stages. It adds the residual image of the early input features to the image acquired in the first stage, then passes to the next stage. Ablation experiments show that the two-stage network has significant benefits for image quality and visual effects. The experimental results on RESIDE, O-Haze, and I-Haze datasets show that the method is superior to advanced methods in dehazing effectiveness. 1 INTRODUCTION Due to the presence of haze, outdoor images often suffer from color distortion, blurring and the quality of other visible captured content, and limit the performance of tracking [1], object detection [2, 3] and other advanced computer vision task [4] performance. Therefore, single image dehazing has been a challenging and highly focused problem to help other computer vision tasks significantly improve prediction accuracy. Many previous dehazing methods are based on the classical atmospheric scattering model [5] : I x = J x t x + A 1 − t x , \begin{align} \begin{split} I{\left(x \right)} = J{\left(x \right)}t{\left(x \right)} + A{\left({1 - t{\left(x \right)}} \right)}, \end{split} \end{align} (1)where I ( x ) $I(x )$ and J ( x ) $J(x )$ donate the hazy image and the clear image at the pixel location x , respectively. A is the global atmospheric light, t ( x ) $t(x )$ is the transmission image. Besides, t ( x ) = e − β d ( x ) $t(x ) = {e^{ - \beta d(x)}}$ , where β is the atmospheric scattering parameter and d ( x ) $d(x)$ is the distance from the camera to the scene. To estimate a clean image from a hazy image, researchers have proposed some priority-based [6-10] methods based on the atmospheric scattering model. These methods attempt to estimate atmospheric channel A ( x ) $A(x )$ and medium transport graph t(x) by using handcrafted priors, such as dark channel [8] and non-local priors[6]. However, these physical priors are not always reliable, and in the absence of priors and constraints, it is difficult to accurately estimate A ( x ) $A(x )$ and t(x). Therefore, incorrect estimation can further reduce the dehazing performance, resulting in artifacts, color distortion and inadequate dehazing . This method has achieved some results, but there is still much room for improvement. In recent years, with the development of deep learning technology, many dehazing methods based on deep learning have been proposed. These methods use convolutional neural networks (CNNs) to extract features and directly learn the mapping between hazy and dehazing image pairs. For example, Cai [11] used an end-to-end model in the early stage to estimate the image transmittance by using a neural network on the atmospheric scattering model, so as to restore a clean image.[12] proposed the patch-map-based hybrid learning DehazeNet, which integrates these two strategies by using a hybrid learning technique involving the patch map and a bi-attentive generative adversarial network.Li [13] did not estimate the atmospheric light separately, and designed a lightweight CNN dehazing network to directly generate clear images. [14] proposed a deep residual haze network, which restored the hazy images by subtracting the learned negative residuals from the hazy images. The modular design could effectively aggregate the context information, and the nonlinear activation function can further improved its presentation ability and convergence speed. [15] proposed the Trident Dehazing Network (TDN), which learns the fog to non-fog image mapping through automatic haze density recognition. [16] modified the basic cGAN formula and combined the conditional generation confrontation network with the codec structure to restore more realistic and clear images. Later, [17] applied a hybrid network of encoder and decoder networks to a task-oriented image defogging network, and further improved the accuracy of the multi-stage defogging algorithm by gradually filtering the haze residuals. In order to obtain clear images from a multi-scale perspective, [18] proposed a multi-scale boosted dehazing network with dense feature fusion based on the U-Net architecture. At the same time, through enhanced and error feedback mechanism to correct spatial information and use non-adjacent hierarchical features. [19] designed a multi-scale deep network to learn the mapping relationship between hazy images and transmission maps. First, the coarse-scale network is used to predict the overall projection image to obtain a rough dehazing image, and then the fine-scale network uses high-frequency detail information to repair and obtain a clear image. [20] proposed a multi-scale topological network (MSTN) to explore features at different scales. At the same time, multi-scale feature fusion module (MFFM) and adaptive feature selection module (AFSM) are designed to realize feature selection and fusion at different scales. [21] proposed a fast multi-scale end-to-end dehazing network called FAMED-Net, which is composed of encoders and fusion modules of different scales to directly learn clean images.All the above are based on supervision, while Li [22] designed a deep convolutional neural network (CNN) with supervised learning branch and unsupervised learning branch in the synthetic data and real images. The two branches constrain the network by supervised loss function and priority-based clean image attribute, respectively. [23] proposed an adaptive model for end-to-end training of image migration and dehazing networks. First, the image translation network is applied to translate images from one domain to another, reducing the difference between the synthetic domain and the real domain. The two images before and after migration are used to train the dehazing network. In addition, researchers have recently made many attempts [24-26]. These deep learning methods have the problem of low efficiency of long-distance feature transmission. Most of them adopt skip connection to improve the efficiency of feature transmission, which limits the ability of network to extract deep features. In addition, existing deep models extract features through the convolution operation of shared weights. The shallow convolution can only focus on local features, but cannot pay attention to global features. If the convolution module is added, the computational consumption will increase, which will bring great burden to the network and limit the processing capacity of the network. Transformer [27] was originally proposed to solve the Seq2Seq problem by being able to process word vectors in parallel, effectively encoding longer sequences of global information, and speeding up model reasoning. [27-29] These Transformer based models have demonstrated strong presentation capabilities and have made breakthroughs in the NLP space. Inspired by the power of Transformer, researchers have recently applied transformers to computer vision (CV) tasks. Chen [30] train a sequence Transformer to auto-regressively predict pixels, and achieved high accuracy in image classification tasks. Dosovitskiy [31] proposed a visual Transformer, which directly applies a pure Transformer to the sequence of image patches, which achieves the most advanced performance on multiple image recognition benchmarks. In addition to image classification, Transformer has been used in multiple task areas such as image representation [32], image segmentation [33], target detection [34, 35], pose estimation [36-38]and pre-training [39]. However, when migrating this model to the visual field, there are still some problems that cannot be ignored, such as the scale difference between NLP and visual target. Using Transformer in CV, the computational complexity is the square of the image scale, which increases the computational complexity. Based on the above problems,[40] designed a hierarchical transformer (Swin Transformer), which introduced the idea of locality to perform self-attention calculations in the window area without overlap, which solved the problem of large image scales and greatly reduced the computational complexity. In this paper, we propose a new end-to-end two-stage dehazing network built with swin transformer feature presentation capabilities. The first stage is the basic feature extraction of foggy image, and the second stage is the local feature extraction and fusion. In addition, we have improved the network level of swin transformer and added an inter-block supervision mechanism (ISM) to refine the characteristics of the previous encoder unit. These features are supervised and selected through the channel attention mechanism, and they are passed to the next decoder unit for interaction and aggregation. This design reduces the loss of feature information between codecs due to upsampling and downsampling, which is conducive to enriching features and better achieving information feature interaction. In addition, we add a fusion attention mechanism(FAM) between the stages. It can add residual images with early input features to the images acquired by the decoder, and pass the attention-guided features to the next stage for processing. FAM is conducive to the transmission of the basic real characteristic signals recovered at each stage, and can effectively avoid the feature transmission with less information to improve the efficiency of network information transmission. In summary, our contributions are as follows: (1) We propose a new two-stage single image dehazing network and design a new transformer-cnn codec, which can effectively extract and fuse local and global features. (2) In order to improve the efficiency of the interaction between encoder and decoder blocks, we propose an inter-block supervision mechanism (ISM). ISM effectively reduces the loss of feature information caused by upsampling and downsampling, and is beneficial to enrich features. (3) Joining the fusion attention mechanism (FAM) between the stages is conducive to the integration of deep and shallow features between the stages, and improves the learning ability of the network. 2 RELATED WORK In this section, we briefly review the development and application of the single-image dehazing algorithm and transformer, and then introduce our approach. Single image dehazing The purpose of single Image dehazing is to predict clean images from blurred images. There are two categories: priority-based methods and learning-based methods. The priority-based method relies on the physical model and restores the image through various prior knowledge or assumptions [6, 8, 10, 41]. Dark channel prior (DCP) is one of the outstanding representatives of priority-based methods. He [8] assumed that each pixel with a value close to zero has at least one color channel, and the imaging model can be used to directly estimate high-quality dehazing images. Zhu [10] proposed a method of restoring image color attenuation by establishing a linear model to estimate the depth of field information. Instead of using an indefinite image priority-based approach, Berman [6, 41] proposed a non-local priors, which predicts atmospheric light by identifying haze lines and estimating the transmission of each pixel. Although these prior-based methods have achieved certain results, these priors are easily violated in practice and are less robust in complex actual scenarios. Unlike the priority-based methods, learning-based methods are data-driven. Early methods [11, 19] used deep neural networks to estimate the transmission map and atmospheric light in physical scattering models. For example, [11] proposed a trainable end-to-end system dehazing network for medium transmission estimation. It uses CNN to estimate the transmission map, and then performs image restoration through traditional stages and atmospheric scattering models. [19] embedded the atmospheric physical scattering model into a multi-scale deep neural network, which was composed of a coarse scale network and a fine scale network: the former predicted the overall transmission pattern of the image, while the latter refined local features. Different from [11, 19], [13] reconstructed the physical scattering model and directly generated clear images through lightweight CNN. However, due to the inaccurate estimation of medium transmission and atmospheric light, these methods will lead to the reconstruction error between the restored image and the clear image. Recently, some methods [18, 42-45] learn the mapping from blur to clear image, estimate the blur image through an end-to-end network and achieve good results. Based on the principle of enhancement and error feedback, [18] designed a multi-scale enhanced dehazing network with dense feature fusion based on U-net architecture. The spatial information is retained by the dense feature fusion module, and the enhanced decoder is used to gradually restore the clean image. GCANet [43] proposed an end-to-end gated context aggregation network. The smoothing expansion technique is used to eliminate the meshing artifacts and the features of different levels are fused by the gating sub-network. FFA-Net [44] combines channel attention with pixel attention mechanism to deal with different types of information flexibly. At the same time, the feature fusion module is designed to learn feature weight adaptively, which gives more weight to important features and promotes the retention and transmission of shallow features of the network. [45] proposed a novel end-to-end network based on attention-based multi-scale estimation, which can effectively alleviate the bottleneck problem often encountered by traditional multi-scale methods and reduce artifacts in the final output.These deep learning methods all have the problem of low efficiency of long-distance feature transmission. Most of them use skip connection to improve the efficiency of feature transmission, but they can only extract limited shallow feature information. In addition, existing models usually achieve highly nonlinear mapping from fuzzy inputs to results by learning a large number of filters, which inevitably increases computational consumption. Therefore, we introduced the swin transformer block in this article to address the above issues. Transformer Transformer [27] is a deep neural network mainly based on self-attention mechanism, which is mainly used in the field of natural language processing. Compared with CNN, when it calculates the correlation between two positions, the number of operations required does not increase with the distance. It can calculate in parallel and effectively extract, process global features, and encode long sequences. Because of its powerful presentation capabilities, researchers have recently applied transformer to computer vision tasks. Chen[30] trained a sequence transformer to automatically regression predict pixels and achieved high accuracy in image classification tasks. In addition, Dosovitskiy [31] proposed a visual transformer (VIT), which can be directly applied to image patch sequences to complete image classification tasks. Transformer is also used for advanced visual tasks such as target detection [34, 46, 47], lane detection [48], and segmentation [33, 49]. DETR [34] uses the convolutional network to learn the two-dimensional representation of the picture, and then uses the encoder to extract the features, which are input to the DECODER together with the position encoding for decoding, finally target detection. Zhang [46] proposed a feature pyramid transformer (FPT) to make full use of feature interactions across spaces and scales, and used self-attention modules to enhance feature fusion of feature pyramid network. Chi [47] proposed bridging visual representation (BVR), in which enhanced features of the primary representation are bridged with information from the auxiliary representation through multi-head attention modules. A single representation can effectively adapt to the final detection performance. Liu [48] combined the Transformer network, convolutional neural network and hungarian loss to design a tiny lane detection framework, and improved its performance by learning the global background using the transformer network. SETR[33] divided the image into a series of patches, and transformer is used to code and model the global context. [49] applied transformer to panoramic segmentation and proposed that Max-Deeplab directly predicted the mask of category markers using a mask translator, and trained the model with panoptic quality inspired loss through binary matching. Liu [40] proposed a hierarchical transformer to solve the problem of large image scale, and used the sliding window method to replace the multi-head self-attention module in transformer. It greatly reduces the computational complexity and makes the network computation increase linearly. This method solves two problems of transformer application in CV field, which are large change of target scale and excessive computation. We apply the transformer block in this method to our single image dehazing network to extract the features of the input image at different resolutions. 3 PROPOSED METHOD Overall architecture The overall architecture of the proposed network is shown in Figure 1, which consists of two stages. In the first stage, the codec structure is combined with swin transformer and CNN to extract global features where the encoder uses swin transformer unit. In order to divide the input image into sequence for input, we first divide the image into 4 × 4 overlapping patches, and the feature dimension of each patch is changed to 4 × 4 × 3 = 48 ${4} \times {4} \times {{\rm 3 = 48}}$ , and then the linear embedding layer is applied to map the image to any dimension (expressed as C) on the original feature. The transformed patch uses several encoder units (composed of swin transformer and patch merging) to model the global relationship and transform the hierarchical features. At the same time, the width and height of the feature map are continuously reduced and the receptive field is expanded. Although the combination of a transformer and raw up sampling can produce certain performance, this strategy is not the best use in the dehazing task. Although swin transformer solves the problem of edge loss caused by patch segmentation, the resolution of transformer image is much smaller than that of original image.It will inevitably lead to the loss of low-level details (such as image object edge shape), and the correlation between features is small. Therefore, in order to compensate for this information loss, we use CNN structure cascaded upsampling as decoder to recover the spatial order and decode the hidden features accurately, so as to predict the full resolution of the final dense output. FIGURE 1Open in figure viewerPowerPoint The architecture and configuration of the proposed method consists of two stages. The first stage is composed of transformer encoder and CNN decoder, which realizes the extraction of global features of the hazy image. The second stage is the enhancement stage, which further extracts and fuses local features to obtain the restored image. In addition, we have added an inter-block supervision mechanism (ISM) between the encoder unit and the decoder unit to refine the features passed in from the previous encoder unit. FAM stands for the inter-stage fusion attention mechanism, which can add residual images with early input features to the images obtained by the decoder and transfer the attention-guided features to the next stage In the second stage, the original resolution image is manipulated to extract local features, and the interaction and fusion of global and local features are carried out. This stage mainly includes several feature units, each of which is composed of local feature extraction module (LFM) and combination feature collection module (CFM). LFM does not use any down-sampling operations and is used to extract the rich global high-resolution features of the image. CFM combines the global features such as color, shape, and texture obtained by the codec with the local features obtained in the second stage, which can more fully and effectively extract information such as image edges, lines and special attributes, thereby effectively achieving image restoration. In addition, we designed an inter-block supervision mechanism (ISM) between the decoder unit and the encoder unit. As shown in Figure 2, ISM first refines the features from the previous encoder unit with a 1 × 1 convolution. Then supervise the selection through the channel attention mechanism, and pass them to the next decoder unit for interaction and aggregation. This design reduces the loss of feature information between encoder and decoder due to upsampling and downsampling, which is beneficial to enrich features and realize information feature interaction better. FIGURE 2Open in figure viewerPowerPoint (a) Inter-block supervision mechanism (ISM). (b) Swin transformer block. (c) Fusion attention mechanism (FAM). Under the supervision of the early input features, the residual images are fused and the attentional guided features are transferred to the next stage Similarly, in order to enhance the connection between the stages, we designed a fusion attention mechanism (FAM) between them. Under the supervision of early input features, FAM added residual images with early input features to the images obtained by the decoder. Finally, the characteristics of attention guidance are transferred to the next stage for processing. It can facilitate the transmission of basic real characteristic signals recovered at each stage. At the same time, it can effectively avoid the feature transmission with less information and improve the efficiency of network information transmission. Encoder-decoder 3.2.1 Transformer encoder Different from transformer used in CV in the past, swin transformer introduces the hierarchical construction method commonly used in cnn, and introduces the locality idea in the multi-head self-attention (MSA) module to perform self-attention calculation in the window area with no overlap. Swin transformer can solve the problem of large image scale through hierarchical transformer. In addition, the sliding window design can avoid edge loss and reduce the complexity of calculation. Therefore, we directly apply swin transformer in the encoder to encode the feature representation from the decomposed patch. As shown in Figure 1, our encoder unit is mainly composed of patch merge and swin transformer block.The maximum or average pooling of downsampling commonly used in CNN often loses information. Therefore, for each encoder unit, we use patch merging to carry out a sliding window operation to reduce the image resolution. With the increase of network depth, the number of patches is constantly reduced to achieve the purpose of layered representation. Then, the patch sequence obtained is mapped to n-dimensional space. Taking the first patch merging as an example, before feature transformation, the original adjacent 2 × 2 patches are joined together, and a linear layer is applied to 4C-dimensional input features. At this point, the number of patches decreases by 2 × 2 = 4 $2 \times 2 = 4$ times, the output dimension is set to 2C, and the feature size is set to H 8 × W 8 ${\textstyle {H \over 8}} \times {\textstyle {W \over 8}}$ from the original H 4 × W 4 ${\textstyle {H \over 4}} \times {\textstyle {W \over 4}}$ . Next, the output features of patch merging enter two continuous swin transformer blocks for characteristic transformation (as shown in Figure 2). Next, the output feature z l − 1 ${z^{l - 1}}$ of Patch Merging enter two continuous Swin Transformer blocks for characteristic transformation (as shown in Figure 2). Swin transformer block consists of a layer specification layer (LN), an MLP, a residual layer, and a window multi-head self-attention layer. W-MSA) and shifted window multi-head self-attention (SW-MSA) across the wire. In order to reduce the complexity and improve the global information exchange capability, this paper adds a computational attention map with cross-block connection in the local area. Therefore, when calculating the self-attention, the relative position offset B is added to the corresponding head,then the l-layer output feature z l − 1 ${z^{l - 1}}$ can be written as follows: z l ′ = W − M S A ( L N ( z l − 1 ) ) + z l − 1 , \begin{align} \begin{split} {z^{l^{\prime }}} = W - MSA(LN({z^{l - 1}})) + {z^{l - 1}}, \end{split} \end{align} (2) z l = M L P ( L N ( z l ′ ) ) + z l ′ , \begin{align} \begin{split} {z^l} = MLP(LN({z^{l^{\prime }}})) + {z^{l^{\prime }}}, \end{split} \end{align} (3) z l + 1 ′ = W − M S A ( L N ( z l ) ) + z l , \begin{align} \begin{split} {z^{l + 1^{\prime }}} = W - MSA(LN({z^l})) + {z^l}, \end{split} \end{align} (4) z l + 1 = M L P ( L N ( z l + 1 ′ ) ) + z l + 1 ′ , \begin{align} \begin{split} {z^{l{{\rm + 1}}}} = MLP(LN({z^{l{{\rm + 1}}^{\prime }}})) + {z^{l{{\rm + 1}}^{\prime }}}, \end{split} \end{align} (5)where z l ${z^{l}}$ and z l − 1 ${z^{l - 1}}$ represent the output of multi-head self-attention, z l ${z^{l}}$ and z l − 1 ${z^{l - 1}}$ represent the output of MLP. In addition, self attention is calculated as follows: A ttention ( Q , K , V ) = SoftMax Q K T d + B V , \begin{align} \begin{split} A{{\rm ttention(Q,K,V)}} = {\rm SoftMax} \left(\frac{{Q{K^T}}}{{\sqrt d }} + {\rm B}\right) {\rm V}, \end{split} \end{align} (6)where Q , K , V ∈ R M 2 × d $Q,K,V \in {R^{{M^2} \times d}}$ denote the query, key, value matrices. M2 and d represent the number of patches in a window and the dimension of the query or key, respectively. And