Towards Robust Rain Removal with Unet++

1  Introduction

Image rain removal is an image preprocessing method that deals with the inverse problem of removing the rain effect from an image, while highlighting its details to meet application-specific requirements and make it more suitable for human-machine recognition [1,2]. It is common to find rain in videos and images taken under bad weather conditions [3,4]. The presence of rain not only negatively affects the visual quality of an image or video, but also reduces the performance of application-specific tasks, such as object segmentation, recognition, tracking, and autonomous driving [5,6]. Therefore, video/image deraining is a crucial preprocessing step and important research area in the field of computer vision [7,8]. In the image rain removal task, many background elements and rain streaks are fused together, making them difficult to distinguish. When a convolutional neural network (CNN) is used to extract the rain trace information, it is easy to accidentally extract and remove the background information. It is worth to mention that although generative adversarial networks (GANs) can reconstruct a more realistic image after rain removal [9,10], performance evaluation reveals that their performance is actually worse than the CNN-based reconstruction method. It can be inferred that many of the detail textures reconstructed by GANs are false and unreal.

In this paper, an innovative rain removal method is proposed, and a first attempt is made to investigate, improve and evaluate the robustness of rain removal methods. The main contributions are summarized as follows:

•   An improved rain removal method is proposed based on Unet++ which achieves good performance in removing rain streaks from images and reducing their disruptive effects.

•   A new loss function that includes Structural Similarity Index (SSIM) loss, edge loss, and L1 loss is adopted to improve the efficiency of the proposed approach for image deraining.

The remainder of this paper is organized as follows. In Section 2, we introduce some related work in this context. The framework of our method is presented in Section 3. We describe the proposed approach in Section 4. Experimental results are presented in Section 5, and conclusions are drawn in Section 6.

2  Related Work

Image deraining is the process of recovering a clear scene from the naturally blurred landscape due to bad weather conditions. Conventional image deraining algorithms generally rely on manual priors [11,12]. When the scene conditions do not meet these priors, such algorithms often produce an unrealistic output, which leads to non-ideal quality in the final restored image. In order to address this issue, with the great success of deep learning in the field of computer vision, numerous researchers have proposed various image rain removal methods based on CNNs [13,14]. A CNN is used to estimate the transmission map or directly predict the clear image [15,16]. These methods have been proved to be efficient and superior to priori-based methods, with significant performance improvement.

Single image rain removal algorithms are generally divided into conventional model-driven methods and data-driven deep learning methods [17,18]. Due to the strong automatic feature learning ability of deep networks, the single image rain removal algorithm based on deep learning has surpassed the conventional model-driven methods in recent years and caught the interest of field researchers [19,20]. To solve the issue of insufficient interpretability and incomplete integration with physical structures inside general rain streaks, Wang et al. proposed a model-driven deep neural network for single image rain removal [21]. Aiming to address real-world rain removal problems, namely the failure in handling real-world rainy scenes, Quan et al. proposed an innovative cascaded network architecture to remove rain streaks and rain drops [22]. It is well known that improving the performance of a single image rain removal algorithm based on deep learning mainly focuses on two aspects: the data quality of the rain map and the design of an image rain removal algorithm [23]. Numerous scholars have proposed new synthetic rain maps or rain map rendering methods to enhance the expression ability of rain map datasets. Wang et al. constructed the relatively real large-scale rain map dataset SPA-data through a semi-automatic method [24]. Ye et al. proposed a new framework which can jointly learn real rain generation and removal procedures [25]. Li et al. proposed a new deep network named REcurrent SE Context Aggregation Net (RESCAN). After decomposing the rain removal process into multiple stages, RESCAN incorporated a recurrent neural network (RNN) to preserve the useful information in previous stages and benefit the rain removal in later stages [26]. Guo et al. proposed a model-free deraining method named EfficientDeRain for the single-image deraining, which is over 80 times faster than the state-of-the-art methods [27]. Li et al. developed a method named ESnet [28], which used layered Long Short-Term memory (LSTM) for recurrent deraining and fine-grained encoder feature refinement. Choose Rain100H as a testing dataset, Table 1 compares the outcomes of several popular neural network methods, including JORDER [9], PreNet [19], RCDnet [20], RESCAN [26], EfficientDeRain [27], and ESnet [28]. It can be seen from Table 1 that PreNet is able to achieve the best result in terms of PSNR, and RCDnet is able to achieve the best result in terms of SSIM. EfficientDeRain can achieve a good performance in terms of both PSNR and SSIM.

3  Framework

The proposed deraining model is based on Unet++ and has two parts: training and optimization, as shown in Fig. 1.

Figure 1: Framework

Unet++ develops from Unet, which is closely similar to Fully Convolutional Networks (FCN) [29]. Unet was proposed shortly after FCN, but both were published in 2015 [30]. In the framework of the proposed approach, Unet++ is used to train the deraining model. With the increase in the number of network layers, the layer by layer downsampling will continuously incur information loss, so the small receptive field feature is used as supplement in order to increase the accuracy of location information. It can be used to achieve better results in rain removal as well. In the optimization step, a new loss function that aggregates three parts–SSIM loss, edge loss and L1 loss–is used in image deraining.

4  Deraining Model

4.1 Pixel Level Extended Filtering Network with Unet++

Similar to Unet, the Unet++ model consists of two parts: left encoder part and right decoder part. The Unet++ architecture is composed of four repetitive structures: convolution, downsampling, upsampling, and skip sampling. The number of characteristic channels for each downsampling is doubled. In the decoder part, similar to the coding layer, deconvolution also consists of four repeating structures. In Unet++, each point in horizontal direction is connected so that different feature levels can be captured [31]. The shallow layer is more sensitive to small targets while the deep layer is more sensitive to large targets, because the receptive fields at different depths have different sensitivities to targets of different sizes. The advantages of those operations can be integrated by concatenating the features.

Unet++ architecture is used to extract the features of the image. The proposed deraining network architecture is shown in Fig. 2. The rainy day image was used as input to estimate the pixel-level kernel G for rain removal, which is similar to the method in [32]. Through off-line training of clean image pairs on rainy days, the kernel prediction network can predict the spatial variable kernels of rain stripes with different thicknesses and intensities, while preserving the object boundary.

Figure 2: Rain removal network architecture

The pixel-level kernel can be estimated by:

G=UPN(I~),(1)

where I~ represents the input rainy image and UPN() denotes a network similar to UNet++. The pixel-level extended filtering network with Unet++ is shown in Fig. 2. Multiple magnified images are used for filtering and fusion. In order to enable the method to deal with multi-scale rain streaks without affecting efficiency, each prediction kernel is extended to three scales through extended convolution.

In order to reduce parameters and time costs, pixel-wise multi-dilated filtering for convolutional layer is suggested in [33,34]. This strategy is also adopted in the proposed approach. When the rain stripe covers a large area of the image, the large-scale kernel is used to effectively remove the rain from the relevant pixels that are far from the rain area. Through Eq. (2), the derained pixel can be predicted:

I^l(g)=∑s,h=g+sGg(s)I~(h),(2)

where g is defined as 2D coordinates for a pixel in the derained image and h is defined as 2D coordinates for a pixel in the rainy image, s is a step parameter which ranges from (−K−12,−K−12) to (K−12,K−12), and l is a dilation scale factor. In the experiments, four scales are set, i.e., l = 1, 2, 3 and 4. At last, a 3 × 3 convolution kernel is used to fuse the four images and obtain the final output.

4.2 Loss Function

Unlike some common methods which only consider one loss function, three loss functions are considered for training the proposed network, namely L1, SSIM and Edge loss functions. Given the derained image I~ and the clean image I as the real image, the proposed loss function is defined as:

ℒ(I~,I)=α‖I~−I‖1−βSSIM(I~,I)−γEdge(I~,I)(3)

where α=0.9, β=0.1 and γ=0.15 were empirically determined.

SSIM is an index to measure similarity of two images. Among the two images compared by SSIM, one is an original image and the other is a distorted image. The SSIM is defined as,

SSIM(X,Y)=(2uXuY+C1)(2ρXY+C2)(uX2+uy2+C1)(ρX2+ρY2+C2),(4)

where X and Y are two input images, uX is the mean value of X, uY is the mean value of Y, ρX is the variance of X, ρY is the variance of Y, ρXY is the covariance between X and Y, while C1 and C2 are two constants.

Typically, the purpose of edge detection is to significantly reduce the data size of the image while retaining the original image attributes. While numerous edge detection algorithms exist, Canny algorithm is very well-known and commonly used.

The Canny edge detection algorithm can be expressed as follows. Given an image I, it is first smoothed using a Gaussian filter. Then, the first-order partial derivative operator is used to find the partial derivative of the image along the horizontal direction Ex and the vertical directionEy. The orientation of the gradient is calculated according to the following formula:

θ=arctan⁡(EyEx),(5)

and the amplitude of the gradient is defined as:

E=Ex2+Ey2.(6)

The non-maximum value of the gradient amplitude is suppressed, i.e., the maximum value of the local gradient is found.

4.3 Algorithm

Our rain augmentation learning algorithm is shown in Algorithm 1.

In Algorithm 1, Dirichlet and Beta distributions are used to preprocess the weights. At each training iteration, we generate a rain map via function Rain_aug and then augment it to the clean or rainy images. The new rainy images are then used for training the network and the fusion convolution layer.

5  Experimental Evaluations

5.1 Evaluation Metrics

In our benchmark, two types of performance evaluation metrics are considered along with three datasets (Rain100H, Rain100L, Rain1400) [26,27]. Common quality measures including peak signal-to-noise ratio (PSNR) and SSIM are used for evaluating results fidelity.

5.2 Comparisons

The proposed approach will now be compared with several popular methods: JORDER [9], PreNet [19], RCDnet [20], RESCAN [26], EfficientDeRain [27] and ESnet [28], which are representative neural network methods. Some examples of deraining for person, landscape, and object images are shown in this section.

Fig. 3a shows three person images and Fig. 3b their rainy versions. One can see from Fig. 3 that the proposed method can recover the details of the human face and remove the rain from the image. Comparison results with regard to PSNR and SSIM are shown in Table 2. The best results are highlighted in bold. One can see from Table 2 that the proposed method obtained the best results.

Figure 3: Qualitative results of different methods on several human images. From left to right: (a) Ground truth images, (b) Rainy image, (c)-(i) Rain removal results by JORDER, PreNet, RCDnet, RESCAN, EfficientDeRain, ESnet and proposed

Fig. 4a shows different landscape images, and their rainy versions are shown in Fig. 4b. One can see from Fig. 4 that the proposed method can avoid mistakes in deraining. Comparison results in terms of PSNR and SSIM are shown in Table 3. The best results are highlighted in bold. One can see from Table 3 that the proposed method obtained the best results.

Figure 4: Qualitative results of different methods on several landscape images in Rain100H testing data set. From left to right: (a) Ground truth images, (b) Rainy image, (c)-(i) Rain removal results by JORDER, PreNet, RCDnet, RESCAN, EfficientDeRain, ESnet and proposed

Fig. 5a shows different object images, and their rainy images are shown in Fig. 5b. One can see from Fig. 5 that, in the proposed method, object characteristics (such as fur and edge) are better preserved after deraining. Comparison results of PSNR and SSIM values are shown in Table 4. The best results are highlighted in bold. One can see from Table 4 that the proposed method obtained the best results too.

Figure 5: Qualitative results of different methods on different objects. From left to right: (a) Ground truth images, (b) Rainy image, (c)-(i) Rain removal results by JORDER, PreNet, RCDnet, RESCAN, EfficientDeRain, ESnet and proposed

Rain100H was chosen as testing dataset. The results obtained by different methods are shown in Table 5. It can be noticed that the proposed method achieves the best results in terms of both PSNR and SSIM.

5.3 Ablation Studies

The loss function has three parameters: α, β and γ. We examine the impact of different strategies proposed in this paper for performance improvement. Further experiments are carried out, in which β and γ are modified while α = 0.9 is fixed. Fig. 6 demonstrates the effect of modifying β∈[0.1, 0.3], while γ = 0.1. PSNR and SSIM results are shown in Fig. 7.

Figure 6: Result showing the effects of varying parameter β, while α=0.9 and γ=0.1

Figure 7: Result showing the effects of varying β

We can know from Fig. 7 that with the increasing noise in the parameter β, which represents the weight of SSIM loss, PSNR and SSIM values decrease. When β = 0.10, our method achieves the best PSNR and SSIM results.

Fig. 8 demonstrates the effect of modifyingγ∈[0.1,0.3], and the corresponding PSNR and SSIM results are shown in Fig. 9.

Figure 8: Result showing the effects of varying parameter γ while α=0.9 and β=0.1

Figure 9: Result showing the effects of varying γ

We can know from Fig. 9 that with the increasing noise in the parameter γ, which represents the weight of edge loss, PSNR and SSIM values fluctuate. When γ=0.15, our method achieves the best PSNR and SSIM results.

Setting too large values for β or γ may lead to poor results. According to our experiments, β=0.1 and γ=0.15 achieve the best results.

6  Conclusions

A kernel prediction network with Unet++ is designed and used to filter rainy images. Then, to optimize the gap between synthetic and real data, a new loss function which consists of three parts, including SSIM loss, edge loss, and L1 loss, is further proposed along with an effective data optimization method that helps to improve the performance in real rainy image handling.

•   Numerous examples of deraining for person, landscape, and object images demonstrated that our method is better than other common methods according to the PSNR and SSIM evaluation metrics.

•   Experiments on datasets show that the proposed method outperforms the state-of-the-art methods under all evaluation metrics. It has improved nearly 1.3% from the conventional methods in terms of the PSNR evaluation metric.

•   We also found that setting too large values for β or γ may lead to poor results. According to our experiments, the parameters in the loss function α=0.9, β=0.1 and γ=0.15 achieve the best results.

Further research based on multi-head-attention in more complex environments and conditions is still necessary. This should be a next research priority.

Funding Statement: This work was supported by National Natural Science Foundation of China (61772179), Hunan Provincial Natural Science Foundation of China (2022JJ50016, 2020JJ4152), the Science and Technology Plan Project of Hunan Province (2016TP1020), Scientific Research Fund of Hunan Provincial Education Department (21B0649), Application-Oriented Characterized Disciplines, Double First-Class University Project of Hunan Province (Xiangjiaotong [2018]469), Discipline Special Research Projects of Hengyang Normal University (Grant No. XKZX21002).

Conflicts of Interest: The authors declare that they have no conflicts of interest to report regarding the present study.