A Light Dual-Task Neural Network for Haze Removal

TL;DR

This paper introduces LDTNet, a lightweight neural network that simultaneously estimates transmission maps and removes haze from images in a single step, outperforming existing methods on synthetic and real-world data.

Contribution

The paper presents a novel dual-task neural network that jointly performs haze removal and transmission map estimation, reducing reliance on artificial priors and improving generalization.

Findings

Achieves superior haze removal performance compared to state-of-the-art methods.

Effectively estimates transmission maps alongside haze removal in a single network.

Demonstrates robustness on both synthetic and real-world images.

Abstract

Single-image dehazing is a challenging problem due to its ill-posed nature. Existing methods rely on a suboptimal two-step approach, where an intermediate product like a depth map is estimated, based on which the haze-free image is subsequently generated using an artificial prior formula. In this paper, we propose a light dual-task Neural Network called LDTNet that restores the haze-free image in one shot. We use transmission map estimation as an auxiliary task to assist the main task, haze removal, in feature extraction and to enhance the generalization of the network. In LDTNet, the haze-free image and the transmission map are produced simultaneously. As a result, the artificial prior is reduced to the smallest extent. Extensive experiments demonstrate that our algorithm achieves superior performance against the state-of-the-art methods on both synthetic and real-world images.

Tables

Table 1. TABLE I: Average PSNR and SSIM of different dehazing methods on the synthetic dataset

Metrics	DCP	CAP	MSCNN	DehazeNet	AOD-Net	LDTNet
PSNR	14.6713	20.9303	20.2855	21.9690	21.3655	24.6156
SSIM	0.8432	0.9452	0.9274	0.9463	0.9419	0.9517

Table 2. TABLE II: Average PSNR and SSIM of dehazing results using ARE, CRE, SRE and NRE

	Metrics	ARE	CRE	SRE	NRE
DehazeNet	PSNR	21.7716	21.8800	22.0891	20.0235
	SSIM	0.9450	0.9423	0.9416	0.3469
MSCNN	PSNR	20.1496	20.1152	20.4302	19.8722
	SSIM	0.9258	0.9240	0.9202	0.4017
AOD-Net	PSNR	21.1600	21.0030	21.3820	19.9497
	SSIM	0.9399	0.9395	0.9355	0.4098
LDTNet	PSNR	24.1181	24.2780	24.6344	22.2765
	SSIM	0.9459	0.9468	0.9441	0.4925

Equations

$$I(x)=J(x)t(x)+A(1-t(x)),$$

$$t(x)=\exp (-\beta d(x)),$$

$$L=(1-\alpha )L_D(J(x),J^{\ast }(x))+\alpha L_T(t(x),t^{\ast }(x)),$$

Full text

I Introduction

Haze is a common phenomenon when the light is absorbed and scattered by the turbid medium. It leads to low visibility and contrast in outdoor scenes. Hazy image can be described using the Atmospheric scattering model, which was firstly proposed in 1976 [1] and widely used by computer vision and graphics community. Later, Narasimhan and Nayar [2, 3, 4, 5] improved the model and re-formulated it as:

[TABLE]

where $x$ denotes the pixel locations in the image, $I(x)$ demotes the hazy image observed, $J(x)$ demotes the real scene without haze, $A$ demotes the atmosphere light, and $t(x)$ demotes the medium transmission. The term $t(x)$ can be further written as:

[TABLE]

where $d(x)$ is the depth of the scene and $\beta$ indicates the scattering coefficient of the atmosphere.

It is difficult to use hazy images directly for most computer vision tasks like object detection [6, 7], tracking [8, 9], pose estimation [10, 11], behavior analysis [12, 13], and search [14, 15]. Researchers have thus been devoting great efforts in haze removal to restore images with high quality, among which haze removal from a single image becomes the focus. For example, Tan [16], Fattal [17] and He [18] implemented single-image dehazing using hand-crafted features, upon which the approaches of [19, 20, 21, 22] were proposed.

All above algorithms, however, rely on specific hand-crafted features, which are not able to fully characterize the hazy images. To this end, recent research has been focused on applying Convolutional Neural Network (CNN) that is able to automatically extract features to handle this task. DehazeNet [23] and MSCNN [24] design CNNs to estimate the transmission map of the hazy image and subsequently use it to estimate atmosphere light. Then, the haze-free image is computed using Eq. (1). In AOD-Net [25], the atmosphere scattering model is re-expressed, and the atmosphere light, scattering, and transmission map are rewritten into one matrix. AOD-Net estimates this matrix and introduces addition layers as well as multiplication layers to compute the re-expressed formula. Although AOD-Net estimates transmission map and atmosphere light at the same time, it first produces an intermediate product, i.e., the matrix of estimated parameters,and then computes the haze-free image by an artificial formula based on the obtained matrix.The errors of the intermediate step may therefore propagate to the dehazing part and thus downgrade the results. GFN [26], on the other hand, directly estimates the clear scenes from hazy images but relies complex pre-processing operations including white balance, contrast enhancing, and gamma correction.

In this paper, we propose a light end-to-end dehazing deep network, termed Light Dual-Task Network (LDTNet). In contrast to prior models like AOD-Net that decouples the process into two steps, ours estimates dehazed images from hazy ones in one shot, in other words, we do not rely on any intermediate output. Furthermore, our model does not rely on artificial priors, such as atmospheric scattering model and the one of Eq. (1). To facilitate the feature learning process, we introduce Multitask Learning (MTL) [27] to estimate the haze-free image and transmission map simultaneously.

The reason we incorporate the transmission map estimation as an auxiliary task is that, the features learned for this task can benefit our main task of dehazing. With a single-task dehazing network, however, it may be very tough to learn such features. Also, learning the main task alone bears the risk of overfitting [28], which can be alleviated by the joint learning with the auxiliary one. In fact, this phenomenon of auxiliary task benefiting the main one has been observed in a wide domain of high-level vision and text tasks [29, 30, 31, 28, 32],

Our contribution is therefore a light Multitask dehazing deep network that does not depend on artificial priors, and that produces haze-free image and transmission map simultaneously in one shot. Our model yields superior results, compared to the state of the art, on both synthetic and real-world data.

II Model

In this section, we introduce our proposed LDTNet. We start by showing the architecture design and then introduce the loss function. Unlike prior models that heavily rely on artificial priors, either the hand-crafted features or hypothetical dehazing models, our approach jointly learns two tasks, the main task of dehazing and auxiliary one of transmission map estimation, without human-provided priors in one shot.

II-A Architecture Design of LDTNet

The proposed LDTNet, with the help of the auxiliary task, is able to restore haze-free image with a lightweight structure. The LDTNet is composed of three cascaded convolutional layers, where the restored image and the estimated transmission map are obtained in the third layer. We show the architecture of LDTNet in Fig. 1.

The two tasks share the first two convolutional layers by hard parameter sharing. The sizes of convolutional kernels in these two layers are all $3\times 3$ and the output channels are 30 and 40 respectively. The input three-channels RGB hazy image is concatenated to these two layers severally as three additional feature maps. This operation provides information contained in the input image, ensuring the refinement of the features layer by layer. There are two parts of the last layer. They combine the feature maps of the second convolutional layer in different ways to reconstruct the haze-free image and transmission map respectively. The sizes of convolutional kernels in these two parts are all $1\times 1$ while the output channels are 3 and 1 respectively.

In LDTNet, no pooling layers are used and zero pixels are padded to the features maps, ensuring the size of the output image to be consistent with that of the input map. Furthermore, batch normalization is applied after the first two convolutional layers. Bilateral Rectified Linear Unit (BRelu) [23] is adopted as our activation function. Specifically, BRelu is a modified version of Rectified Linear Unit (Relu) [33] with the upper limit set to be 1. It ensures that the pixels in the restored image are constrained in the range of $[0,1]$.

II-B Loss Function Design of LDTNet

In LDTNet, we simultaneously tackle dehazing and transmission estimation. We take the loss function to be

[TABLE]

where $L_{D}$ and $L_{T}$ correspond to the dehazing loss and transmission loss respectively, and $\alpha$ balances the two. In our implementation, both $L_{D}$ and $L_{T}$ take the form of square loss, defined on the pixel-wise difference between the ground truth and the prediction of the network.

III Experiments

In this section, we show our experimental validation of the proposed LDTNet. We first introduce the baselines methods we use for comparison, and then present our training strategy. Afterwards, we show the effectiveness of our dual-task learning, followed by the comparative results on synthetic and real-world test images. We finally provide the robustness analysis of different methods.

III-A Baselines

We compare our proposed LDTNet with several state-of-the-art methods briefly introduced as follows.

• •

DCP [17]: The thickness of haze is estimated first and then the haze-free image is recovered using Eq. (1).

• •

CAP [34]: Color attenuation prior is utilized for estimating the scene depth, which is further used for computing the haze-free image.

• •

MSCNN [24]: The transmission map is estimated and refined using two CNNs of different scales, which is then used to obtain the haze-free image by Eq. (1).

• •

DehazeNet [23]: The transmission map is estimated using a CNN with a novel nonlinear activation function.

• •

AOD-Net [25]: The transmission map and the airlight are jointly learned using a CNN.

III-B Training

To train the LDTNet, we synthesize 10,000 triples of hazy images, haze-free images and transmission maps, all of which are resized to $240\times 320$ pixels based on the NYU depth [35] dataset using (1) and (2). We set $A\in(0.7,1.0)$ and $\beta\in(0.5,1.5)$, and thus our dataset covers various atmosphere situations, multifarious levels of haze, as well as different weather conditions. For LDTNet, we use Adam [36] to be our optimizer and set the batch size to be 4. We implement our model using Tensorflow [37] and the Tensorlayer [38] package. It takes about 17 hours to train the LDTNet for 100 epochs on a Titan X GPU with Intel i7 CPU.

III-C Validity of Dual-task Learning

To validate our dual-task learning, we compute the Mean Square Error (MSE) between our result and ground truth under various values of $\alpha$ on our validation set, as shown in Fig. 2. When $\alpha=0$, the auxiliary task is removed and only the main task dehazing plays a role, in which case we obtain the largest MSE. This shows that our multitask learning is indeed helpful. Also, as can be seen, the performance of network comes to the best when $\alpha=0.4$ on our validation set.

III-D Performance on Synthetic Test Dataset

We use a synthetic dataset including 21 pairs of stereo images generated using the Middlebury stereo database [39, 40, 41] to verify the performance of different models. The atmosphere light $A$ is set to be 0.85 and the scattering coefficient $\beta$ is set to 1. These values correspond to the medians of the domains of $A\in(0.7,1)$ and $\beta\in(0.5,1.5)$. It takes about 0.3 second to produce a dehazed image during testing.

We show the obtained mean PSNR and SSIM of the results in Table I. LDTNet achieves the highest PSNR and SSIM scores. In Fig. 3, we show the comparative dehazing results of five examples: Midd, Cloth, Bowling,Aloe and Monopoly.

As can be seen, LDTNet yields visually plausible results which are very similar to the ground truths. The restored images of the other methods have larger color distortions or over-saturations. They are over-sensitive to regions of light colors like white, as they appear similar to haze. For instance, in the case of Cloth, baseline methods tend to produce yellowish colors for the regions that are supposed be of white color.

Baseline methods fail to produce very accurate results, in part due to their hand-crafted features or their two-step nature. In the cases of DCP and CAP, they rely on a limited set of handcraft features that may not be expressive enough for some scenes. MSCNN and DehazeNet, on the other hand, first estimate the intermediate transmission maps and then produce the dehazed results. The errors occurred in the transmission estimation may thus propagate to dehazing and negatively influence the results. Although AOD-Net estimates the transmission map and atmosphere together using a CNN, it still relies on an artificial formula to obtain the haze-free image.

By contrast, LDTNet directly learns a mapping from hazy images to the haze-free ones by taking the transmission map estimation as an auxiliary task. As a result, LDTNet does not suffer from the limitations of intermediate products and artificial priors.

III-E Performance on Real-world Test Images

We qualitatively compare our algorithm with DCP, CAP, MSCNN, DehazeNet and AOD-Net on 50 challenging real-world images. We show a few of them in Fig. 4. As can be seen from the second column, DCP leads to over enhancement especially on the first three images. CAP is yields better results than DCP does in terms of over enhancement, but fails to preserve textural details in the region of similar colors, like the mountains in the second image and the hair of the girl in the third. MSCNN, DehazeNet, and AOD-Net also suffer a certain degree of over enhancement. In addition, MSCNN causes hue distortion in the last image of the fifth column. The results of LDTNet are indeed more visually plausible, without noticeable color distortions or loss of details.

III-F Robustness Analysis

We conduct robustness analysis, as done in DehazeNet [23], for baselines methods and ours using four types of evaluations, i.e., airlight robustness evaluation (ARE), coefficient robustness evaluation (CRE), scale robustness evaluation (SRE) and noise robustness evaluation (NRE).

In Table IV, we show the mean PSNR and SSIM on the Middlebury stereo dataset. In ARE, we synthesize 210 hazy images with $\beta=1$ and $A\in(0.7,1.0)$. Similarly, the same number hazy images are synthesized with $A=0.85$ and $\beta\in(0.5,1.5)$ for CRE. To analyze the influence of the scale variation, we select four scale coefficients, i.e., $1,0.8,0.6,0.4$, to generate different scale images with $\beta=1$ and $A=0.85$. Finally, we add three types of noises to the hazy images generated with $\beta=1$ and $A=0.85$ for NRE. The three kinds of noises are Gaussian noise, Poisson noise and salt & pepper noise.

As can been seen in Table II, we achieve the best performance in the four types of evaluations. This is because our dual-task learning approach improves the network’s capability of task-relevant features extraction and generalization. As a result, LDTNet can extract effective features from the hazy images even in the presence of noise and various $A$ and $\beta$. Also, since we fuse the original hazy images into the convolutional layers, the information of the source image can be preserved and utilized to a greater extent.

IV Conclusion

In this paper, we have presented a light dual-task neural network, LDTNet, which takes as input hazy images and produces dehazed ones in one shot, without any intermediate results. The auxiliary task, transmission map estimation, proves to be helpful for enhancing dehazing and for improving the network’s generalization capability. We conduct quantitative and qualitative evaluations, and compare our results with those of the state-of-the-art methods on both both synthetic and real-world hazy images. LDTNet yields the most promising results in terms of both accuracy and robustness. In our future work, we will explore to simultaneously conduct dehazing and other tasks, such as high-level object detection [42] and tracking [43, 44] ones, and low-level super-resolution [45] and image restoration [46] ones, where the both tasks could potentially benefit each other.