Dehaze-UNet: A Lightweight Network Based on UNet for Single-Image Dehazing

Abstract

Numerous extant image dehazing methods based on learning improve performance by increasing the depth or width, the size of the convolution kernel, or using the Transformer structure. However, this will inevitably introduce many parameters and increase the computational overhead. Therefore, we propose a lightweight dehazing framework: Dehaze-UNet, which has excellent dehazing performance and very low computational overhead to be suitable for terminal deployment. To allow Dehaze-UNet to aggregate the features of haze, we design a LAYER module. This module mainly aggregates the haze features of different hazy images through the batch normalization layer, so that Dehaze-UNet can pay more attention to haze. Furthermore, we revisit the use of the physical model in the network. We design an ASMFUN module to operate the feature map of the network, allowing the network to better understand the generation and removal of haze and learn prior knowledge to improve the network’s generalization to real hazy scenes. Extensive experimental results indicate that the lightweight Dehaze-UNet outperforms state-of-the-art methods, especially for hazy images of real scenes.

1. Introduction

Images captured in bad weather will be disturbed by impurities such as haze, which will reduce the contrast of the image and the visibility of details, which will affect the subsequent computer vision task. Therefore, image dehazing has received widespread attention from academia and industry. The atmospheric scattering model (ASM) [1] is widely used to describe the imaging process of hazy scenes:

$$\mathit{I}\left(x\right)=\mathit{J}\left(x\right)\mathit{t}\left(x\right)+\mathit{A}(1-\mathit{t}(x\left)\right),$$

(1)

where $\mathit{I}\left(x\right)$ represents the image collected in hazy weather, $\mathit{J}\left(x\right)$ denotes the clear image, A represents the global atmospheric light, and $\mathit{t}\left(x\right)$ represents the transmission map.

Restoring $\mathit{J}\left(x\right)$ from known only $\mathit{I}\left(x\right)$ is clearly an ill-posed problem. The earlier method is to estimate the parameters in ASM to restore $\mathit{J}\left(x\right)$ through statistical priors. Some scholars have proposed some famous priors. He et al. [2] proposed the well-known dark channel prior (DCP). Yan et al. [3] proposed the bright channel prior. Other priors include the nonlocal prior [4], gradient profile prior [5], etc. This prior knowledge has strongly promoted the research of image dehazing algorithms, and many works based on this prior knowledge have achieved better results [6,7,8]. Although these prior-based methods take into account the degradation mechanism of hazy images, when the prior assumption is not established, the effect of dehazing will be correspondingly worse. Recently, learning-based dehazing methods have received extensive attention from researchers. Cai et al. [9] estimated parameters through an end-to-end convolutional neural network (CNN) and then restored the clear image through ASM. Ren et al. [10] designed a dehazing network model to estimate parameters. Li et al. [11] improved ASM and proposed AODNet. In recent years, the relatively novel dehazing work has included end-to-end methods [12,13,14,15], attention mechanism methods [16,17], weakly supervised methods [18], semisupervised methods [19,20,21], knowledge-distillation-based methods [22,23], contrastive-learning-based methods [24,25], and Transformer-based methods [26,27,28,29]. Although these methods have been very successful, these dehazing network models are often accompanied by complex network structures and large model parameters, which have high computational overhead. This greatly affects the operating efficiency of the model, and it is difficult to deploy them on terminal devices with limited computing resources and fast response requirements.

In this work, considering the practical application of image dehazing, we propose a lightweight dehazing model: Dehaze-UNet. Based on U-Net [30], Dehaze-UNet can extract multiscale features, which significantly reduces computational overhead while ensuring dehazing performance, making it more suitable for terminal devices, especially for devices with limited computing. In Dehaze-UNet, we design two novel components, the LAYER module and the ASMFUN module. The LAYER module mainly aggregates the haze features of different hazy images through the batch normalization (BN) layer [31], so that the network focuses more on haze to ensure improved dehazing performance. The ASMFUN module incorporates ASM to allow Dehaze-UNet to better comprehend haze generation and removal. In addition, the ASMFUN module allows Dehaze-UNet to learn prior knowledge under the constraints of the physical model so that the network can thoroughly analyze and understand the characteristics of haze during training. We highlight the key contributions of this work as follows:

• We propose a novel end-to-end lightweight dehazing framework that has few parameters and very small computational overhead but achieves excellent dehazing performance.
• To allow the network to better aggregate the features of haze, the LAYER module is designed. This module mainly aggregates the haze features of different hazy images through the batch normalization layer so that the network can focus on the hazy scenes in the image.
• We revisit the use of a physical model in the network. We designed the ASMFUN module and embedded it into the network’s intermediate layer to allow the network to gain a deeper insight into the mechanisms of haze generation and removal. This module combines the physical model to allow Dehaze-UNet to learn prior knowledge during training and improve its generalization.

2. Related Works

2.1. Prior-Based Methods

According to ASM [1], the image dehazing process is an underconstrained problem. Some traditional methods usually achieve image dehazing through prior knowledge. Through observation, He et al. [2] found that in most nonsky local image patches, there are some pixels with an intensity close to zero, so they proposed the famous dark channel prior (DCP) and estimated the unknown quantities in ASM based on this prior. However, as DCP pointed out, this prior does not hold true in the sky area of the image, so the haze-free image it recovered is prone to problems such as halos and color casts in the sky area. Zhu et al. [32] proposed a color attenuation prior and estimated the scene depth through a supervised learning method. However, a simple linear model is not enough to describe complex real hazy scenes, so this method makes it difficult to achieve an obvious dehazing effect for hazy scenes with large concentrations. Berman et al. [33] proposed that color clusters become different haze lines, through which the hazy image can be restored. Li et al. [34] proposed a method to fuse DCP and BCP, which effectively solved the problem that DCP is invalid in the sky area. These prior-based methods have had some success. However, manually designed priors are not able to obtain sufficient image statistical information, meaning prior-based dehazing methods have limitations in many cases. This will lead to unsatisfactory dehazing effects, such as artifacts.

2.2. Deep-Learning-Based Methods

DehazeNet, proposed by Cai et al. [9], and MSCNN, proposed by Ren et al. [10], were the first to use deep learning to estimate parameters and restore clear images through ASM. Liu et al. [16] proposed GridDehazeNet, which includes a preprocessing module and a postprocessing module. The network does not rely on ASM and learns through multiscale network estimation. Qin et al. [17] also introduced an attention mechanism in their FFANet, which includes channel attention and spatial attention. FFANet simultaneously considers different pixel information with uneven haze distribution. However, the processing time of a single image is as long as 3.55 seconds, which is not conducive to real-time applications of the dehazing system. Dong et al. [12] proposed MSBDN with dense feature fusion. AECR-Net [24] applies contrast loss to the image dehazing. Yu et al. [35] proposed a model based on frequency–spatial domain dual guidance to explore and extract haze-related features. In addition, some research also attempts to enhance the generalization of dehazing models through unsupervised or semisupervised methods. Li et al. [36] proposed an unsupervised dehazing method YOLY. Since there is no need to train a deep model on hazy and clean image pairs, this method effectively alleviates the domain transfer problem. Zhao et al. [18] proposed RefineDNet, which does not rely on paired hazy–nonhazy image data sets during the training process. Chen et al. [19] proposed PSD guided by physical priors to reformulate the real-world dehazing task into a generalization framework from synthetic to real. The model is pretrained on synthetic datasets and fine-tuned with unsupervised learning based on real-world datasets. Yang et al. [20] proposed D4, and the training process of this model only requires unpaired hazy and hazy-free images, but it can recover the atmospheric scattering coefficient.

3. Proposed Method

3.1. Motivation: Dehaze-UNet

The main purpose of dehazing is to provide preprocessing services for high-level machine vision tasks, and it should appear in a way that a small model assists other large model tasks. Therefore, we propose a lightweight dehazing network model, Dehaze-UNet. In order to make the network model small enough, Dehaze-UNet refers to the classic U-Net structure. The U-Net structure can generate feature maps of different sizes through downsampling at different scales so that multiscale features can be extracted and fused. Dehaze-UNet has fewer network layers, so the number of network parameters is smaller, and the computational overhead is very low. Different from the previous U-Net and its variant structures, Dehaze-UNet as an image dehazing task contains two novel modules: the LAYER module and the ASMFUN module.

3.2. Overall Architecture

Figure 1 shows the overall architecture of Dehaze-UNet, which can be considered a variant of the U-Net network. From the input $\mathit{I}$ to the output $\mathit{J}$, the layers of the network are as follows: 3 × 3 convolution (i.e., 3 × 3 CONV in Figure 1), downsampling (i.e., DOWN in Figure 1), 1 × 1 convolution (i.e., 1 × 1 CONV in Figure 1), LAYER, downsampling, 1 × 1 convolution, LAYER, LAYER, ASMFUN, LAYER, upsampling (i.e., UP in Figure 1), ASMFUN, LAYER, upsampling, and 1 × 1 convolution. The 3 × 3 convolution is to extract shallow features and increase the number of channels, and the 1 × 1 convolution is to fuse channel features given an image pair ($\mathit{I}$,${\mathit{J}}_{GT}$), where $\mathit{I}$ is the hazy input and ${\mathit{J}}_{GT}$ denotes the corresponding ground truth (GT) image. Dehaze-UNet is trained by a simple L1 loss, as shown in Equation (2), where i denotes the i-th training sample, and N represents the total number of samples participating in training. Although Dehaze-UNet only employs a simple L1 loss, it can still achieve excellent performance.

$$\zeta =\frac{1}{N}\sum _{i=1}^N{∥{\mathit{J}}^i-{\mathit{J}}_{GT}^i∥}_1,$$

(2)

The downsampling module of Dehaze-UNet is implemented by a convolutional layer with size 2 and stride 2. The upsampling module of Dehaze-UNet is realized by a 1 × 1 convolution layer and a pix-shuffle [37] layer.

3.3. LAYER Module

The designed LAYER module is shown in Figure 2. Each layer is mainly a batch normalization (i.e., BATCHNORM in Figure 2), convolution (i.e., CONV in Figure 2), ReLU activation function (i.e., ReLU in Figure 2), etc. Inspired by the work of batch normalization [31], the proposed LAYER module is a parallel structure, which extracts features separately and finally fuses them in the form of a point product. Through the batch normalization layer, the network can more effectively aggregate the haze features of different hazy images and focus more intently on the haze scene so that the subsequent ASMFUN module can better fit the physical model.

3.4. ASMFUN Module

The designed ASMFUN module is shown in Figure 3. First, it is a grouped dilated convolution, and the convolution kernel size is 5 × 5. The aim of employing dilated convolution is to enhance the receptive field, and the grouped convolution is to extract the global features of a single channel as much as possible and to make the features of each channel unaffected. Then, after the batch normalization layer, it is activated by the ReLU function, and then there are two parallel structures. This parallel structure is inspired by SKNet [38] and is used to perform high-dimensional abstraction and estimation of $\mathit{t}\left(x\right)$ and A in ASM. These two parallel structures are composed of an adaptive average pooling layer (i.e., AdaptiveAvgPool2d in Figure 3), 1 × 1 convolution, ReLU activation, and 1 × 1 convolution. The adaptive average pooling layer can effectively extract features, and can change the size of the feature map to a value of 1 × 1 through adaptive parameters; after 1 × 1 convolution, ReLU activation, 1 × 1 convolution, a multichannel 1 × 1 feature map is obtained, which represents the estimated $\mathit{t}\left(x\right)$ and A with high-dimensional features, respectively. Please note that $\mathit{t}\left(x\right)$ and A obtained here are not equivalent to the two-dimensional $\mathit{t}\left(x\right)$ and one-dimensional A in the standard ASM but an abstract result from a higher-dimensional perspective, and then the feature map of the output result is obtained through ASM. The ASM here is not the end-to-end addition of the ASM in most other dehazing methods to obtain hazy or hazy-free images but is embedded as an intermediate layer of the network to operate on the intermediate feature map. This is to allow the network to better understand the formation and removal of haze during the training process. This idea is inspired by the related theory of physics-informed machine learning (PIML) [39]. According to ASM (Equation (1)), $\mathit{J}$ can be expressed as

$$\mathit{J}\left(x\right)=\left(\mathit{I}\right(x)+A(\mathit{t}\left(x\right)-1\left)\right)/\mathit{t}\left(x\right),$$

(3)

However, in the ASMFUN module, ASM is embedded in the intermediate layer to operate on the feature map, not on the final actual output image. Therefore, in order to prevent the situation where the denominator in Equation (3) is 0, a change was made to Equation (3). Assuming that the input of the ASMFUN module is INPUT, and the highly abstract $\mathit{t}\left(x\right)$ and A obtained by the two parallel structures are $\tilde{\mathit{t}}\left(x\right)$ and $\tilde{\mathit{A}}$, respectively, the feature map $OUT$ output by the module can be expressed as

$$OUT=(INPUT+\tilde{\mathit{A}}(\tilde{\mathit{t}}\left(\mathit{x}\right)-\mathit{1}))·\tilde{\mathit{t}}\left(x\right),$$

(4)

The change from the division of Equation (3) to the multiplication in Equation (4) does not mean that the ASM is violated. This is a technical change based on the mechanism of deep learning and experimental observations, which still allows Dehaze-UNet to learn prior knowledge under the constraints of the physical model so that the network can better understand the characteristics of haze. In addition, it is found through experiments that the network trained using the original ASM is difficult to converge, and the loss value appears infinite, while the network trained using the improved Equation (4) is relatively stable.

4. Experiments

4.1. Implementation Details

We performed a lot of comparative experiments on synthetic hazy images and real-world hazy images, including quantitative and qualitative comparative experiments, as well as an ablation study. Our experimental equipment is equipped with 4 NVIDIA GeForce GTX 1080Ti GPUs and uses the Adam optimizer and cosine annealing strategy to train Dehaze-UNet. In order to obtain better model parameters, a pretraining method is adopted; that is, the initial model parameter values are obtained through pretraining with a batch size of 256, a learning rate of 0.002, and a block size of 256 × 256. Formal training is then performed with a batch size of 32, a learning rate of 0.002, and a block size of 512 × 512, where the value of the learning rate is reduced to 0.00002 by a cosine annealing strategy. To verify the performance of Dehaze-UNet, 9 advanced dehazing methods were selected for qualitative and quantitative experimental comparisons. These methods include DCP [2], DehazeNet [9], AODNet [11], GridDehazeNet [16], MSBDN [12], FFANet [17], RefineDNet [18], PSD [19], D4 [20], and DehazeFormer [29]. Among them, DCP [2] is a prior-based method, and the rest are based on deep learning methods. We adopt the well-known PSNR and SSIM indicators [40] for experiments. The higher the scores of these two indicators, the better the recovery performance of the method.

4.1.1. Synthesize Image Dataset

We chose the well-known RESIDE dataset [41], trained on ITS and OTS and then tested on SOTS.

4.1.2. Real-World Image Dataset

For real-world hazy images, the O-Haze [42], NH-Haze [43], and Dense-Haze [44] datasets were chosen. Dehaze-UNet is retrained with the NH-Haze dataset and then evaluated on O-HAZE; similarly, Dehaze-UNet is retrained with the O-Haze dataset and evaluated on NH-Haze and Dense-Haze.

4.2. Experiments on Synthetic Hazy Images

The SOTS dataset from RESIDE [41] was selected for comparative experiments. The SOTS-Indoor is a synthetic indoor hazy image dataset. Two images from this dataset were randomly selected for qualitative comparison experiments. Figure 4 shows the results of clear images restored by different dehazing methods. From Image 1, compared with the ground truth (GT), the image restored by DCP is darker. The color information of the images restored by DehazeNet, AODNet, GridDehazeNet, MSBDN, D4, and other methods is insufficient, and some of the detailed information marked by the red box in Figure 4 is also biased. The image restored by the RefineDNet method has obvious noise, especially the part marked by the red box, and the image restored by the PSD method is obviously brighter. The results restored by the Dehaze-UNet method and the FFANet method are similar and are closest to the GT image. From the PNSR/SSIM index score results of the restored image, the PNSR index score of the image restored by the Dehaze-UNet method is the highest, indicating its excellent performance. The SSIM index score of the Dehaze-UNet method is second only to FFANet, and the difference is only 0.0034, indicating that the performance of the two methods is close. The recovery in Image 2 has a similar situation. From the score, the PSNR index score of the Dehaze-UNet method is second only to FFANet, and the SSIM index score is lower than the FFANet and GridDehazeNet methods. Although the index score restored by Dehaze-UNet did not achieve the highest value, it still shows that Dehaze-UNet is competitive, and it further demonstrates that the FFANet exhibits superior performance on the synthetic dataset.

The SOTS-Outdoor is a synthetic outdoor hazy image dataset. Two images from this dataset were randomly selected for qualitative comparison experiments. Figure 5 shows the visual results of images restored by different dehazing methods. As can be seen from Image 3 in Figure 5, the image restored by DCP is darker and has severe distortion in the sky area. The image restored by the DehazeNet and AODNet methods is insufficiently protected in the sky area. The images restored by the GridDehazeNet, MSBDN, and D4 methods have chromatic aberrations in the sunlight peripheral part enlarged by the red box. The brightness of the images restored by the FFANet and PSD methods exceeds the GT image. The image restored by the RefineDNet method has severe noise; while the image restored by the Dehaze-UNet is closest to the GT image. From the PNSR/SSIM index results of the restored image, the Dehaze-UNet method has the highest PNSR index, and the SSIM index is second only to GridDehazeNet, with a difference of only 0.0014, which strongly illustrates the advanced performance of Dehaze-UNet. Similarly, from the restored result of Image 4, the restored result of the Dehaze-UNet method is closest to GT in terms of the dehazing effect and the restoration of the color and edge information of the image. Furthermore, the PNSR/SSIM index of the Dehaze-UNet method is the highest, and the PSNR index far exceeds other methods, which illustrates the advanced performance of the Dehaze-UNet.

In addition, in order to further compare the performance of various methods, we conducted quantitative experimental comparisons on all images of the SOTS-Indoor and SOTS-Outdoor datasets. The recovery results obtained by different methods were tested on all images of the two datasets, and the average PNSR/SSIM index score was calculated. From

Table 1. Quantitative experimental comparison results of synthetic images. Bold numbers indicate the highest indicator score, and underlined numbers indicate the next highest score.

Methods	SOTS-Indoor		SOTS-Outdoor
	PSNR	SSIM	PSNR	SSIM
DCP(CVPR 2009) [2]	18.42	0.8259	14.80	0.8029
DehazeNet (TIP 2016) [9]	20.84	0.8122	18.69	0.7426
AODNet (ICCV 2017) [11]	18.73	0.8268	19.64	0.8920
GridDehazeNet (ICCV 2019) [16]	29.86	0.9742	30.34	0.9797
MSBDN (CVPR 2020) [12]	31.98	0.9796	33.49	0.9835
FFANet (AAAI 2020) [17]	32.07	0.9794	30.12	0.9721
RefineDNet (TIP 2021) [18]	20.32	0.8364	20.77	0.9093
PSD (CVPR 2021) [19]	12.46	0.6802	15.31	0.7621
D4 (CVPR 2022) [20]	19.66	0.8563	25.06	0.9396
DehazeFormer-T (TIP 2023) [29]	29.03	0.9390	29.30	0.9711
Dehaze-UNet (Ours)	32.09	0.9421	34.93	0.9799

, the SOTS-Indoor dataset with the Dehaze-UNet method obtained the highest PSNR score, while the SSIM score differed from the highest value by about 0.03. For the SOTS-Outdoor dataset, the Dehaze-UNet achieved the highest PNSR score and the second-highest SSIM score. The Dehaze-UNet method achieved the highest PSNR score, which strongly illustrates its superiority. However, the Dehaze-UNet method did not achieve the highest SSIM score. This is because SSIM mainly evaluates based on the structural similarity of images. Since Dehaze-UNet adopts the U-Net structure, the original structural information of the image may be destroyed during the downsampling process. This is also a point that requires further research and improvement in the future.

4.3. Experiments on Real-world Scenes

Two images were randomly selected from the O-Haze [42], NH-Haze [43], and Dense-Haze [44] datasets, respectively, for comparison, as shown in Figure 6, Figure 7 and Figure 8. From Figure 6, the results recovered by DCP, AODNet, GridDehazeNet, and RefineDNet are darker, and the results recovered by DehazeNet, MSBDN, FFANet, PSD, and D4 contain obvious residual haze. There is basically no obvious haze in the recovery results of Dehaze-UNet. From the PSNR/SSIM indicators, in the recovery results of Image 5, the Dehaze-UNet method achieved the highest score. For the recovery results of Image 6, the Dehaze-UNet achieved the highest PSNR score, and SSIM was second only to the RefineDNet. This also illustrates the competitive performance of the Dehaze-UNet. However, the color of the result restored by the Dehaze-UNet method is distorted, and the color information cannot be maintained while removing haze. This is also a shortcoming of Dehaze-UNet. Future work will improve this shortcoming.

Figure 7 is the experimental comparison result of NH-Haze. As can be seen from Image 7, the results recovered by DCP, AODNet, GridDehazeNet, and RefineDNet are darker. Except for DCP, the remaining methods have very poor dehazing effects, and the recovered images still have a large amount of haze, while the results of the Dehaze-UNet have relatively minimal residual haze. As can be seen from Image 8, the results restored by the DCP and Dehaze-UNet still have the least residual haze, which illustrates the superior performance for outdoor nonuniform hazy images. It also shows that the effectiveness of traditional prior knowledge such as DCP for nonuniform hazy scenes basically exceeds other methods based on deep learning. However, the DCP recovery result is still darker, while the Dehaze-UNet method recovery result is closest to the GT image. It can be seen from the PSNR/SSIM index in the figure that for Image 7, the SSIM score obtained by the Dehaze-UNet is slightly lower than PSD. However, it is obvious that the results of PSD recovery are inferior to the results of the Dehaze-UNet method in terms of color and dehazing ability. This also shows the limitations of the quantitative evaluation index from another perspective; that is, a high PSNR/SSIM index score does not necessarily imply that the visual effect will be equally superior. This is also an issue that requires attention in the field of image dehazing.

In Figure 8, due to the dense haze scene of the Dense-Haze dataset, the recovery results of various advanced dehazing methods are unsatisfactory. However, it can be seen that the recovery result of Dehaze-UNet is still closest to the GT image. From the PSNR/SSIM index in the figure, the index score obtained by the Dehaze-UNet method is the highest, which also verifies the effectiveness of the Dehaze-UNet method in dense haze scenes.

In addition, quantitative experimental comparisons were made on the entire images from the three datasets. The restoration results from different methods were evaluated on all images in these datasets, and the average PSNR/SSIM indicator scores were calculated. The quantitative evaluation results are shown in

Table 2. Quantitative experimental comparison results of the real-world hazy images. Bold numbers indicate the highest indicator score, and underlined numbers indicate the next highest score.

Methods	O-HAZE		NH-HAZE		Dense-HAZE
	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
DCP(CVPR 2009) [2]	14.42	0.5028	11.57	0.4183	12.03	0.4033
DehazeNet (TIP 2016) [9]	15.48	0.6016	11.85	0.4435	11.04	0.4334
AODNet (ICCV 2017) [11]	15.09	0.5433	11.87	0.4246	12.69	0.4591
GridDehazeNet (ICCV 2019) [16]	15.96	0.6395	10.53	0.3874	10.07	0.3986
MSBDN (CVPR 2020) [12]	16.32	0.6366	12.45	0.4903	11.02	0.4533
FFANet (AAAI 2020) [17]	14.68	0.5990	11.93	0.4708	10.35	0.4412
RefineDNet (TIP 2021) [18]	17.32	0.6699	12.79	0.5146	12.91	0.4584
PSD (CVPR 2021) [19]	11.39	0.6095	10.32	0.5274	09.73	0.4345
D4 (CVPR 2022) [20]	16.75	0.6578	12.66	0.5072	11.49	0.4536
DehazeFormer-T (TIP 2023) [29]	15.92	0.6309	12.05	0.4951	10.79	0.4422
Dehaze-UNet (Ours)	21.44	0.6396	17.50	0.5494	16.34	0.4695

. From the PSNR indicator, Dehaze-UNet basically obtained the highest scores and was significantly better than the second-ranked method. Only for the O-Haze dataset, the SSIM result of the Dehaze-UNet method is about 0.03 less than the highest. The above qualitative and quantitative experimental comparisons strongly prove the excellent performance of the Dehaze-UNet method on real scenes.

However, Dehaze-UNet also exhibits limitations in real hazy scenes. For example, it does not maintain the color information of the image well while achieving haze removal, such as Image 5 and Image 6. In nonuniform and dense haze scenes, Dehaze-UNet’s dehazing is still not thorough enough, such as in Image 7 and Image 9. This is mainly because there are too few real image samples participating in the training, and the performance of the trained model is not good enough. These shortcomings are a future research direction.

4.4. Parameter and Computational Complexity Analysis

From

Table 3. Comparison of parameters (# Param) and the number of floating-point operations (FLOPs). Where the FLOPs are tested on the image of size 256 × 256.

Methods	DehazeNet [9]	AODNet [11]	GridDehaze Net [16]	MSBDN [12]	FFANet [17]	RefineDNet [18]	PSD [19]	D4 [20]	DehazeFormer-T [29]	Dehaze-UNet (Ours)
# Param (M)	0.01	0.002	0.956	31.35	4.456	65.8	33.11	10.7	0.686	0.34
FLOPs (GMac)	0.513	0.115	21.49	41.54	288.3	75.41	182.5	2.246	6.658	1.84

, it can be seen that the number of parameters of the Dehaze-UNet model is significantly lower compared with the more famous works since 2019, including GridDehazeNet, MSBDN, FFANet, RefineDNet, PSD, and D4. In addition, the FLOPs of the Dehaze-UNet model are also very small, which are also much lower than GridDehazeNet, MSBDN, FFANet, RefineDNet, PSD, D4, and DehazeFormer-T. This shows that the Dehaze-UNet is a highly lightweight model and is easy to deploy on terminal devices. This is also the key point of our study of this model. Therefore, the proposed method provides a theoretical reference for the lightweight dehazing model.

4.5. Ablation Study

This section mainly studies the impact of Dehaze-UNet’s LAYER module and ASMFUN module on network performance through ablation experiments. In Dehaze-UNet, only the LAYER module, only the ASMFUN module, and both the LAYER and ASMFUN modules were removed to retrain the network and test its performance. Retraining was performed on the ITS dataset and testing on the SOTS-Indoor dataset, respectively. The visual results of the ablation experiment are shown in Figure 9, where the blue text indicates the PSNR/SSIM value. It can be seen from Figure 9 that the image restored without the LAYER and ASMFUN modules has poor recovery of details and is not smooth enough (Figure 9b); the image restored without the LAYER module is smoother (Figure 9c), but the recovery of detail information is still not good; and the image recovered without the ASMFUN module has color distortion (Figure 9d). The result restored by the complete Dehaze-UNet model is closer to the GT image (Figure 9e). The PSNR and SSIM indicator values in the figure also demonstrate the performance of adding these modules.

The quantitative ablation study results are shown in

Table 4. Ablation study of the network, where w/o means to remove the module.

Metrics	SOTS-Indoor
	PSNR	SSIM
w/o LAYER module	22.08	0.7484
w/o ASMFUN module
w/o LAYER module	27.04	0.8879
w/o ASMFUN module	27.68	0.9130
Ours Full Model	32.09	0.9421

. It can be found that adding each module can improve the PSNR and SSIM scores, which illustrates the effectiveness of these two modules.

5. Conclusions and Future Work

This study introduced a lightweight dehazing model, Dehaze-UNet, focusing on its practical application. This model is based on the U-Net structure, which is a pure convolution operation model, and it is designed with a LAYER module and an ASMFUN module. The LAYER module aggregates haze and clear scene characteristics from various hazy images through batch normalization, focusing the network on haze and clear scene features. The ASMFUN module primarily integrates physical models to enable the network to learn prior knowledge during training, enhancing its understanding of haze generation and removal processes and improving overall generalization. Experimental results show that the proposed network model can not only achieve better dehazing effects but also has a small number of model parameters and low computational overhead, making it easy to deploy on terminal devices and improving the practicality of the model.

However, the SSIM index value of the proposed model is not the highest in the experiment, indicating that the model may not preserve the structural information of the image adequately. This is due to the downsampling of the U-Net structure not maintaining enough image structural information. To a certain extent, it is also related to the simple L1 loss function used. Future work should study how to effectively maintain the structural information of the image and try to introduce other loss functions, such as the SSIM loss function, GAN loss function, etc., and appropriately increase the network depth, add convolutional layers, etc., to further improve the model’s performance. Furthermore, the proposed method once again illustrates the necessity of letting the network learn prior knowledge from the fused physical model. In future work, we will try to extend it to video dehazing, addressing dynamic scenes, or exploring domain adaptation for varying atmospheric conditions.