TSNet:A Two-stage Network for Image Dehazing with Multi-scale Fusion and Adaptive Learning

The internal structure of the multi-scale fusion module is shown in Fig. 3. It contains a multi-scale parallel large convolution kernel module (MSPLCK) and an implicit frequency feature enhancement module (IFFE). The MSPLCK acquires multi-scale receptive fields, which provide a better understanding of the global context and structure of the scene 25 . The IFFE obtains features at different frequencies and then integrates these features to reduce the differences in frequency information between inputs and learning objectives.

First, we use BatchNorm to normalize the input F, preventing gradient vanishing and gradient exploding during training, and reducing TSNet’s sensitivity to initial values 37 .

Where Conv_i×i represents a convolution operation with a kernel size of i×i.

Then the MSPLCK is employed with dilated convolution kernels of different sizes to obtain multi-scale receptive fields 25 . It allows the network to better handle global structures. Concat operation is used to concatenate multiple tensors along a specified dimension. Subsequently, an MLP filters and retains more useful features. The outputs of MLP are then added to M(F) to reduce information loss during propagation.

Where $\mathrm{M}(\mathrm{\textbf{F}})$ is k·F, representing assigning a weight k to feature F. The value of k can be dynamically adjusted to an appropriate value during the training process ($0\leq\textit{k}\leq 1$).

Then the IFFE is applied to reduce the differences in frequency information between the inputs and the learning objectives. Finally, the outputs are added to M(F_b) to obtain the output F*, as shown in Eq. 4.

The internal structure of the IFFE is shown in Fig. 3(b). In 38 , it was found that convolutional layers can serve as high-pass filters, reducing low-frequency information in features and increasing the variance and amplitude of high-frequency information. It has an amplifying effect on the high-frequency information. Therefore, in the IFFE, the input feature map I is passed through three convolutional layers sequentially to obtain outputs with different frequencies. Subsequently, the outputs of each convolutional layer are added together to obtain I₁ which contains all diversified frequency information obtained 26 ; 27 . The number of convolutional layers is not fixed and it can be adjusted according to the actual situation. In this paper, three convolutional layers are selected by considering performance and the number of parameters, as shown in Eq. 5.

Then channel attention (CA) 19 is utilized to adjust the weights of each channel and spatial attention (SA) 44 is used in parallel to enhance the perception of critical information. Afterward, by concatenating their outputs, we can obtain beneficial frequency information I₂, as shown in Eq. 6. GAP is global average pooling, and GMP is global max pooling 43 . Average pooling and max pooling can be used to aggregate the spatial information of feature maps 48 .

Finally, a convolutional layer reduces the dimensionality of I₂ and then its output is added to M(I) to get the final output O, as in Eq. 7

In most convolution operations, the shape of the sampling point is fixed, as shown in Fig 4(a). This situation typically damages the texture details of objects, resulting in excessively smooth artifacts and making it challenging for networks to learn useful features in real-world scenes 24 .

The introduction of deformable convolution 31 can effectively address this issue. Compared to regular convolution, deformable convolution learns the offset of the convolution kernel, so that the sampling points of the original regular convolution kernel are shifted, as shown in Fig 4(b). The offset can be dynamically adjusted during the training process to better fit the object itself. Therefore, we introduce the deformable convolution network (DCN) 31 into the ALM to enhance the ability to restore texture details.

The structure of the ALM is shown in Fig. 5. The input X undergoes DCN to obtain X₁, and then a channel attention module is used to adjust the weight of each channel in X₁. Finally, the output of CA is added to M(X) to obtain output Y, as in Eq. 8.

The ALM are deployed before the first upsampling in both the first-stage and second-stage networks. This allows the ALM to understand the preceding features comprehensively while providing a well-prepared input for the subsequent network. In each stage, two ALM are utilized for recovering textures more efficiently while considering the number of parameters.

Unlike previous single-stage dehazing networks, we have added a network for optimizing dehazed images after the dehazing network, serving as the second-stage network of TSNet, as shown in Fig. 2. We observed that after the hazy image passes through the first-stage network, the initial haze has been mostly removed. However, some areas still exhibit noticeable artifacts and color distortion. So we proposed to add the second-stage network to optimize the dehazed images obtained from the first-stage network.

To keep TSNet as concise as possible, the second-stage network is simplified based on the first-stage network, and its structure can be considered a reduced version of the first-stage network. We introduced a loss function between the two stages. The hazy image undergoes processing by the first-stage network, resulting in a dehazed image. Then the dehazed image is utilized as the input for the second-stage network, which is responsible for optimizing the dehazed image. The two stages of TSNet are trained independently, and their weights are saved separately.

A ground truth image can be obtained by adding a hazy image to an opposite fog map, as shown by the dashed box in Fig. 1. Previous networks directly learned ground truth images and output dehazed images. However, the differences between hazy and ground truth images are diverse, including saturation, sharpness, contrast, and so on. So it is challenging for networks to fully consider these differences during the learning process of ground truth images.

Unlike previous methods, we change the learning objective of the first-stage network of TSNet from a ground truth image to an opposite fog image. TSNet outputs a learned opposite fog map b, which is then added to the hazy image to obtain the dehazed image c. This change allows the network to learn the differences between hazy images and ground truth images directly, which enhances its ability to learn the features of various haze types and improves the efficiency of network learning. We abbreviate the above method as CL. The second-stage network optimizes dehazed images by directly outputting the final dehazed images.

Assuming J is the ground truth image in the dataset, and I is the corresponding hazy image. F₁(I) is the output of the first-stage network in TSNet, and F₂(I) is the output of the second-stage network. Then the dehazed image of the first-stage network can be obtained by adding F₁(I) to I.

The loss function $\mathcal{L}_{\text{1}}$ in the first stage consists of $\mathcal{L}_{\text{a}}$ and contrast loss $\mathcal{L}_{\text{b}}$:

Where $\beta$ is a hyperparameter used to balance the contrastive learning loss and $\mathrm{smooth\,L_{1}}$ 45 , $\mathrm{L_{1}}$(x, y) is the L₁ loss, $\omega$_i is the weight coefficient, R_i (i = 1,2, …,n) represents extracting features from the fixed pre-trained ResNet-152 40 model at the i-th layer. Then the loss function $\mathcal{L}_{\text{1}}$ can be expressed as:

The loss function $\mathcal{L}_{\text{2}}$ for the second-stage network can be expressed as:

We evaluated our method on publicly available datasets RESIDE-IN, RESIDE-6K 33 , and Haze-4K 34 , as well as NH-Haze 35 and Dense-Haze 36 . RSIDE-IN contains 13390 pairs of synthetic indoor images for training and the test set is taken from 500 pairs of indoor images in SOTS. RESIDE-6K and Haze-4K are mixed datasets containing both indoor and outdoor synthetic images. Specifically, RESIDE-6K contains 3,000 pairs of indoor images and 3,000 pairs of outdoor images, with 5,000 pairs used for training and 1,000 pairs for testing. Haze-4k contains 4,000 pairs of synthetic images, with 3,000 pairs used for training and 1,000 pairs for testing. NH-Haze and Dense-Haze datasets consist of hazy images generated by a haze machine and their corresponding ground truth images. NH-Haze represents an uneven haze dataset, whereas Dense-Haze represents a dense haze dataset.

We trained our network using one NVIDIA RTX-4090. The optimizer utilized AdamW 41 , with exponential decay rates $\beta$₁ and $\beta$₂ set to 0.9 and 0.999 respectively. The learning rate used the cosine annealing strategy 42 to reduce the initial learning rate from 4×10^-4 to 4×10^-2 during training. As the input and output dimensions of the multi-scale fusion module are the same, it can be used continuously at the same position in TSNet. Therefore, we provided two variants, TSNet-S and TSNet-L, with specific settings as shown in Table 1. The weight $\beta$ of the contrastive loss was set to 0.1. We set $\mathrm{L_{1}}$(x, y) in Eq. 10 after the latent features of the 11th, 35th, and 143rd layers from the fixed pre-trained ResNet-152, with corresponding coefficients $\omega$_i (i = 1,2,3) were set to 1, 1/4, and 1/8.

The size and number of images on different datasets are different, so the settings were different on each dataset 39 , as shown in Table 2.

We perform both quantitative and qualitative analyses on synthetic and real-world datasets, respectively.

Synthetic datasets: In Table 3, we summarize the results on the three synthetic datasets. The data in bold represents the best performance on the corresponding metrics. The performance of TSNet improves as the number of the multi-scale fusion modules increases, as shown by the results of the two variants, TSNet-S and TSNet-L. Compared to most previous methods, TSNet-S has fewer parameters and computational costs. However, its performance is better. TSNet-L achieved the best performance on Haze-4K with 34.95dB PSNR and 0.991 SSIM compared to previous methods. Compared to the second-best-performing method, PMNet 23 , TSNet-L showed a 4.5% increase in PSNR while reducing the number of parameters by 76.9%. Additionally, TSNet-L achieved the highest PSNR (42.60dB) on RESIDE-IN and the highest PSNR (31.31dB) on RESIDE-6K. Its SSIM (0.995) on RESIDE-IN is very close to the highest SSIM (0.996), and SSIM (0.975) on RESIDE-6K is very close to the highest SSIM (0.977). TSNet-S is suitable for tasks with higher demands for real-time performance, while TSNet-L can be applied to tasks with higher requirements for clarity.

Fig. 6 illustrates the dehazed images of TSNet and other models on the RESIDE-IN dataset. Some regions show varying degrees of detail loss and severe blocky artifacts in the images restored by AOD-Net 13 , GridDehazeNet 28 , and FFA-Net 19 . AOD-Net has severe colour distortion. Fig. 7 illustrates the outputs on the RESIDE-6K dataset. AOD-Net, GridDehazeNet, and FFA-Net fail to effectively remove distant haze and exhibit color discrepancies compared to ground truth images. In contrast, TSNet effectively removes haze with fewer artifacts on both RESIDE-IN and RESIDE-6K datasets. Additionally, its color is closer to ground truth images.

Real-world datasets: Many previous state-of-the-art methods performed well on synthetic datasets but exhibited poor performance on real-world datasets. However, TSNet achieved satisfactory performance on real-world datasets as well. From the results in Table 4, TSNet-L achieves the best performance on both Dense-Haze and NH-Haze datasets compared to previous methods. On the Dense-Haze dataset, TSNet-L significantly outperforms the other methods, with 17.48dB PSNR and 0.639dB SSIM. The results demonstrate that TSNet can efficiently remove uneven haze and dense haze.

On the Dense-Haze dataset, the output images of AOD-Net 13 , DehazeFormer-M 21 , and MixDehazeNet-B 25 all exhibit severe artifacts, while the outputs of TSNet appear smoother and better at recovering object details, as shown in Fig. 8. Fig. 9 shows the output images on the NH-HAZE dataset. Many details are lost in the AOD-Net outputs, and residual haze remains in the DehazeFormer-M and MixDehazeNet-B outputs. In contrast, the outputs of TSNet effectively remove haze and restore details such as tree leaves and grass more comprehensively.

To demonstrate the effectiveness of the proposed methods, we conducted ablation experiments in this section. We performed the ablation experiments on the TSNet-S module with the following settings:

Base is the first-stage network of TSNet-S, but it does not include the ALM and the learning objective is ground truth images. (a) involves changing the learning objective from ground truth images to opposite fog maps in Base (CL). (b) adds the ALM to Base. (c) incorporates the second-stage network of TSNet-S into Base (ISSN), and the learning objective of the second-stage network is ground truth images. (d) represents changing the learning objective of the first-stage network (CL) while adding the ALM. (e) uses three methods simultaneously, which represents the complete TSNet-S. To ensure the rigor of the results, we conducted ablation experiments on two datasets, RESIDE-IN, and RESIDE-6K, with all networks trained in the same way.

In Table 5, it can be concluded that the Base, which only adds the multi-scale fusion module, has already surpassed most of the previous state-of-the-art methods. Each of the method (a), (b), and (c) contributes to the improvement of the results, and the dehazing effect improves accordingly when using two or three methods simultaneously.

From Fig. 10, it can be observed that the haze has been greatly eliminated in the outputs of the method (d), but there is residual haze and some color distortion in certain areas. However, the residual haze is further removed and the texture details are better restored in the outputs of the method (e). In addition, the outputs of the method (e) are closer in color to the ground truth images. (d) represents the first-stage network of TSNet-S, and (e) represents the complete TSNet-S. Therefore, the results demonstrate the second-stage network achieves optimization over the outputs of the first-stage network.

The method (c) is to add the second-stage network of TSNet-S to Base, so the original network of Base is deepened. Typically, as the network becomes deeper, the amount of parameters also increase. Therefore, the dehazing effect may also be further improved.

We conducted additional experiments to demonstrate that the improvement in results from our two-stage network design was not solely due to the deeper network. Specifically, we connected the two stages of the network without setting the loss function in the middle, making it a whole network named TS-all. We trained it in the same way as before. Base* represents Base with the addition of methods (a) and (b).

Table 6 shows that the method of two-stage network design is more helpful for improving results on synthetic datasets or real-world datasets compared to TS-all. Therefore, this paper believes that when the network reaches a certain depth, adding loss functions at appropriate positions can regulate the learning process, enabling the network to have clear learning objectives at specific stages.

In addition, the ALM is easy to deploy because of its simple structure, so we conducted experiments to apply it to other networks. Table 7 shows that the ALM can be utilized as a universal module in other networks, enhancing the final dehazing performance of those networks.