Multi-Scale Progressive Fusion Network for Single Image Deraining
Kui Jiang, Zhongyuan Wang, Peng Yi, Chen Chen, Baojin Huang, Yimin Luo, Jiayi Ma, Junjun Jiang
Introduction
Due to substantial degradation of the image content in rain images and videos, traditional image enhancement algorithms struggle to make desirable improvements on image quality. Therefore, developing specialized solutions for image deraining is imperative to a wide range of tasks , e.g. object detection and semantic segmentation.
Traditional deraining methods use simple linear-mapping transformations and are not robust to variations of the input , e.g., rain streaks with various directions, densities and sizes. Recently, deep-learning based methods which operate with convolutional and non-linear layers have witnessed remarkable advantages over traditional methods. Despite obvious improvements on feature representation brought by those methods , their single-scale frameworks can hardly capture the inherent correlations of rain streaks across scales.
The repetitive samples of rain streaks in a rain image as well as its multi-scale versions (multi-scale pyramid images) may carry complementary information (e.g. similar appearance) to characterize target rain streaks. As illustrated in Fig. 1, the rain streaks (highlighted in the white box) in the original rain image share the similar rain patterns with the rain streaks (highlighted in the cyan, pink and dark blue boxes) at different positions as well as those (highlighted in the red, yellow, orange and green boxes) in the 1/2 scale rain image. Therefore, rain streaks both from the same scale (solid arrows) and across different scales (dashed arrows) encode complementary or redundant information for feature representation, which would help deraining in the original image. This correlation of image contents across scales has been successfully applied to other computer vision tasks . Recently, authors in construct pyramid frameworks to exploit the multi-scale knowledge for deraining. Unfortunately, those exploitations fail to make full use of the correlations of multi-scale rain streaks (although restricted to a fixed scale-factor of 2 ). For example, Fu et al. decompose the rain image into different pyramid levels based on its resolution, and then individually solve the restoration sub-problems at the specific scale space through several parallel sub-networks. Such decomposition strategy is the basic idea of many recurrent deraining frameworks . Unlike completing the deraining task from each individual resolution level, Zheng et al. present a density-specific optimization for rain streak removal in a coarse-to-fine fashion, and gradually produce the rain-free image stage-by-stage . However, there are no direct communications of the inter-level features across cascaded pyramid layers except for the final outputs, thus failing to take all-rounded advantages of the correlated information of rain streaks across different scales. Consequently, these methods are still far from producing the desirable deraining results with the limited exploitation and utilization of multi-scale rain information.
To address these limitations of the prior works, we explore the multi-scale representation from input image scales and deep neural network representations in a unified framework, and propose a multi-scale progressive fusion network (MSPFN) to exploit the correlated information of rain streaks across scales for single image deraining. Specifically, we first generate the Gaussian pyramid rain images using Gaussian kernels to down-sample the original rain image in sequence. A coarse-fusion module (CFM) (§3.1) is designed to capture the global texture information from these multi-scale rain images through recurrent calculation (Conv-LSTM), thus enabling the network to cooperatively represent the target rain streak using similar counterparts from global feature space. Meanwhile, the representation of the high-resolution pyramid layer is guided by previous outputs as well as all low-resolution pyramid layers. A fine-fusion module (FFM) (§3.2) is followed to further integrate these correlated information from different scales. By using the channel attention mechanism, the network not only discriminatively learns the scale-specific knowledge from all preceding pyramid layers, but also reduces the feature redundancy effectively. Moreover, multiple FFMs can be cascaded to form a progressive multi-scale fusion. Finally, a reconstruction module (RM) is appended to aggregate the coarse and fine rain information extracted respectively from CFM and FFM for learning the residual rain image, which is the approximation of real rain streak distribution. The overall framework is outlined in Fig. 2. The main contributions of this paper are as follows:
We uncover the correlations of rain streaks in an image and propose a novel multi-scale progressive fusion network (MSPFN) which collaboratively represents rain streaks from multiple scales via the pyramid representation.
To better characterize rain streaks of different scales, we devise three basic modules, coarse-fusion module (CFM), fine-fusion module (FFM) and reconstruction module (RM), to effectively extract and integrate the multi-scale information. In these modules, the complementary information of similar patterns with rain streaks, both within the same scale or across different scales (pyramid layers), is progressively fused to characterize the rain streaks distribution in a collaborative/cooperative manner.
Apart from achieving the state-of-the-art deraining performance in terms of the conventional quantitative measurements (e.g. PSNR and SSIM), we build several synthetic rain datasets based on COCO and BDD datasets for joint image deraining, detection and segmentation tasks. To the best of our knowledge, we are the first to apply mainstream vision-oriented tasks (detection and segmentation) for comprehensively evaluating the deraining performance.
Related Work
In the last few years, substantial improvements have been observed on rain image restoration. In this work, we mainly focus on single image deraining because it is more challenging.
Previous traditional methods for single image deraining fail under the complex rain conditions and produce degraded image contents due to the limited linear-mapping transformation. Very recently, deep-learning based approaches have emerged for rain streak removal and demonstrated impressive restoration performance. For example, Fu et al. introduce a three-layer convolutional neural network (CNN) to estimate and remove rain streaks from its rain-contaminated counterpart. To better represent rain streaks, Zhang et al. take the rain density into account and present a multi-task CNN for joint rain density estimation and deraining. Later, Zhang et al. further incorporate quantitative, visual and discriminative performance into the objective function, and propose a conditional generative adversarial network for rain streak removal. In order to alleviate the learning difficulty, recurrent frameworks are designed to remove rain streaks in a stage-wise manner.
2 Multi-scale Learning
Rain streaks in the air show the apparent self-similarity, both within the same scale or across different scales, which makes it possible to exploit the correlated information across scales for rain streak representation. However, most existing deraining methods ignore the underlying correlations of rain streaks across different scales. Only a few attempts have been made to exploit the multi-scale knowledge. Fu et al. decompose the restoration task into multiple subproblems and employ a set of parallel subnetworks to individually estimate the rain information in a specific pyramid scale space. However, it does not exploit and utilize the correlated information among these pyramid layers. Different from the parallel pyramid framework in , Zheng et al. propose the cascaded pyramid network, which is similar to LapSRN , to iteratively remove rain streaks. However, only the high-level features are used to help the adjacent pyramid representation, which results in losing some useful hierarchical and scale features in a deep cascaded network. The significance of these features produced at different stages has been verified on image reconstruction tasks .
Different from these methods , in this work we introduce a novel framework MSPFN to achieve the collaborative representation of rain streaks across different scales, where the rich multi-scale rain information extracted from the Gaussian pyramid images is progressively aggregated along the pyramid layers and stages of the network. As a result, our predicted rain streak distribution is more accurate via the multi-scale collaborative representation.
Proposed Method
Fig. 2 shows the overall pipeline of our proposed multi-scale progressive fusion network (MSPFN) for image deraining by excavating and exploiting the inherent correlations of rain streaks across different scales. We present the details of each building block and the loss function in the following.
For a given rain image, our method first generates the Gaussian pyramid rain images using Gaussian kernels to down-sample the original rain image into different scales, e.g. 1/2 and 1/4. The network takes as input the pyramid rain images and extracts the shallow features through multiple parallel initial convolution layers (see the first block of “initial layer” in Fig. 2). Based on the initial features from each scale, the coarse-fusion module (CFM) then performs the deep extraction and fusion of multi-scale rain information through several parallel residual recurrent units (RRU), as shown in Fig. 3. The reasons for designing CFM are three folds: (a) To exploit the repetition of rain streaks under the same scale, we apply the recurrent calculation and residual learning to capture the global texture information, making it possible to cooperatively represent target rain streaks. More accurately, we introduce Conv-LSTM to model the information flow of context textures at spatial dimension with the recursive memory, where the contextual texture correlations are transformed into structured cyclic dependencies to capture the complementary or redundant rain information (e.g. the solid arrows in Fig. 1). (b) The multi-scale structure provides an alternative solution to greatly increase the receptive filed to cover more contents while maintaining a shallow depth. (c) The high-resolution representations benefit from the outputs of previous stages as well as all low-resolution pyramid layers via iterative sampling and fusion.
2 Multi-scale Fine Fusion
The outputs of CFM go through the fine-fusion module (FFM) to refine the correlated information from different scales. As shown in Fig. 2, FFM enjoys the similar multi-scale structure with CFM for convenience. Unlike CFM, we introduce the channel attention unit (CAU) to enhance the discriminative learning ability of the network through focusing on the most informative scale-specific knowledge, making the cooperative representation more efficient. To alleviate the computation burden, we apply the strided convolution to reduce the spatial dimension of features, and finally utilize the deconvolution layer to increase the resolution to avoid losing resolution information, resulting in the U-shaped residual attention block (URAB). As depicted in Fig. 4, URAB is composed of several CAUs, along with the short skip connections to help the fine representation of multi-scale rain information. Moreover, long skip connections are used between cascaded FFMs to achieve progressive fusion of multi-scale rain information as well as to facilitate the effective backward propagation of the gradient.
3 Rain Streak Reconstruction
To learn the final residual rain image, we further integrate both low- and high-level multi-scale features respectively from CFM and FFM via a reconstruction module (RM), schematically depicted in Fig. 2. Specifically, the outputs from CFM are concatenated with the outputs from the last FFM, and then a convolution layer is used to learn the channel interdependence and rescale the feature values from the two modules. Similarly, the iterative sampling and fusion of rain information across different pyramid layers are implemented to estimate the residual rain image.
4 Loss Function
Mean squared error (MSE) is the commonly used loss to train the network . However, it usually produces blurry and over-smoothed visual effect with the loss of high-frequency textures due to the squared penalty. In this work, we perform the successive approximation to the real rain streak distribution with the guidance of the Charbonnier penalty function , which is more tolerant of small errors and holds better convergence during training. The function is expressed as
In Equation (1), denotes the predicted residual rain image. The predicted rain-free image is generated by subtracting from its rain-contaminated counterpart . The penalty coefficient is empirically set to .
In order to further improve the fidelity and authenticity of high-frequency details while removing rain streaks, we propose the additional edge loss to constrain the high-frequency components between the ground truth and the predicted rain-free image . The edge loss is defined as
In Equation (2), and denote the edge maps respectively extracted from and via the Laplacian operator . Then, the total loss function is given by
where the weight parameter is empirically set to 0.05 to balance the loss terms.
Experiments and Discussions
We conduct extensive experiments on several synthetic and real-world rain image datasets to evaluate the restoration performance of our proposed MSPFN as well as six state-of-the-art deraining methods. These representative methods include DerainNet , RESCAN , DIDMDN , UMRL , SEMI and PreNet . There is no unified training datasets for all competing methods in this paper, e.g. PreNet refers to JORDER and uses 1254 pairs for training. UMRL refers to and uses 12700 images for training. Therefore, directly taking the results from their papers is unfair and meaningless. To this end, we collect about clean/rain image pairs from for training our network as well as other competing methods for a fair comparison. In particular, these competing methods are retrained in the experiments with their publicly released codes and follow their original settings under the unified training dataset. Separately, the detailed descriptions of the used datasets are tabulated in Table 1. In order to quantitatively evaluate the restoration quality, we adopt the commonly used evaluation metrics, such as Peak Signal to Noise Ratio (PSNR), Feature Similarity (FSIM) , and Structural Similarity (SSIM) .
In our baseline, the pyramid levels are set to 3, i.e. the original scale, 1/2 scale and 1/4 scale. In CFM, the filter numbers of each recurrent Conv-LSTM are respectively set to 32, 64, and 128, corresponding to the gradually increasing resolution. The depths/numbers of FFM (M) and CAU (N) are set to 10 and 3, respectively. We use Adam optimizer with batch size of 8 for training on one NVIDIA Titan Xp GPU. The learning rate is initialized to and reduced by half at every 20000 steps till . We train the network for 30 epochs with the above settings.
2 Ablation Studies
Validation on Basic Components. Using our baseline model (, ), we design six comparison models to analyze the effects of the proposed basic modules (CFM and FFM), multi-scale pyramid framework, and multi-scale progressive fusion scheme on deraining performance. Quantitative results on Test100 dataset are listed in Table 2. From the results, our baseline MSPFN exhibits great superiority over its incomplete versions, including Model1 (single-scale framework with only the original input), Model2 (removing CFM from MSPFN), and Model3 (removing all FFMs from MSPFN), surpassing them by 0.73dB, 0.28dB, and 3.60dB (PSNR), respectively. Moreover, we construct Model4 by applying the fusion strategy in to verify the effectiveness of the proposed multi-scale progressive fusion scheme. It is evident that MSPFN gains a significant improvement over Model4 by 0.54dB with an acceptable complexity increase. Model5 (, ) and Model6 (, ) are the simplified variants of MSPFN with smaller depths. When compared with the single-scale framework (Model1), Model5 has the approximately equal amount of parameters but achieves faster inference speed with the multi-scale pyramid framework. Model6 has the similar computation complexity but more parameters as compared with Model1. The results show that Model5 achieves the comparable performance while it’s a quarter more efficient. Model6 gains the better scores over Model1 by 0.32dB while keeping the similar computation complexity. We attribute these advantages to the effective cooperative representation of rain streaks among different pyramid layers and stages of the network.
Parameter Analysis on and . We assess the influence of the depth of FFM (M) and the number of CAU (N) on deraining performance. Based on our baseline (, ), we construct three comparison models, i.e. MSPFNM17N1, MSPFNM13N2 and MSPFNM8N5, while keeping approximately the same number of parameters. As shown in Table 3, the performance declines with the reduction of M. This indicates the important role of FFM for exploiting the multi-scale rain information in a progressive fashion. When increasing the number of CAU (MSPFNM17N2), it yields a slight improvement (0.13dB), but with additional of the parameters. We also add two models MSPFNM30N1 and MSPFNM5N1 for comparison. The former is designed to pursue a better deraining performance with more FFMs to enhance multi-scale fusion, while the latter is a lightweight model with smaller depth (, ) and width (all filter channels 32). Meanwhile, the strided convolution and deconvolution are employed twice in our proposed U-shaped residual attention block (URAB) of MSPFNM5N1 to further alleviate the computation burden. As we expected, MSPFNM30N1 achieves the best scores for all the metrics. MSPFNM5N1 still obtains the acceptable performance, although being a much lighter network. Considering the tradeoff between efficiency and deraining performance, we set M and N to 17 and 1 respectively in the following experiments.
3 Comparisons with State-of-the-arts
We compare our MSPFN () with other six top-performing deraining methods on five synthetic datasets. Quantitative results are shown in Table 4. One can see that MSPFN achieves remarkable improvements over these state-of-the-art methods. For example, MSPFN surpasses DerainNet and DIDMDN by 9.01dB and 2.74dB, respectively, in terms of PSNR on Test1200 dataset. Visual results on different rain conditions (diverse rain streak orientations and magnitudes) are presented in Fig. 5. MSPFN exhibits impressive restoration performance on all scenarios, generating results with rich and credible image textures while removing main rain streaks. For other comparison methods, they tend to blur the image contents, or still leave some visible rain streaks. For example, only our MSPFN restores the clear and credible image details in the “Giraffe” image, while the competing methods fail to remove rain streaks and their results have obvious color distortion.
3.2 Real-world Data
We conduct additional comparisons on three real-world datasets, including Real200 , Rain in Driving (RID) and Rain in Surveillance (RIS) datasets , to further verify the generalization capability of MSPFN. RID and RIS cover 2495 and 2348 samples, collected from car-mounted cameras and networked traffic surveillance cameras in rainy days respectively. Moreover, we use another two quantitative indicators, Naturalness Image Quality Evaluator (NIQE) and Spatial-Spectral Entropy-based Quality (SSEQ) , to quantitatively evaluate the reference-free restoration performance. The smaller scores of SSEQ and NIQE indicate better perceptual quality and clearer contents. The results are listed in Table 5. As expected, our proposed MSPFN has the best average scores on 200 real-world samples, outperforming the state-of-the-art deraining methods by a large margin. Moreover, we show four representative deraining examples in Fig. 6 for visual comparison. In the last image, obvious rain streaks are observed in the results of other deraining methods, but our MSPFN can well preserve more realistic and credible image details while effectively removing main rain streaks.
3.3 Other Applications
Image deraining under complex weather conditions can be considered as an effective enhancement of image content. It can potentially be incorporated into other high-level vision systems for applications such as object detection and segmentation. This motivates us to investigate the effect of restoration performance on the accuracy of object detection and segmentation based on some popular algorithms, e.g. YOLOv3 , Mask R-CNN , and RefineNet . To this end, we randomly select a total of 850 samples from COCO and BDD datasets to create three new synthetic rain datasets COCO350 (for detection), BDD350 (for detection), and BDD150 (for segmentation) through Photoshop. These rain images are of diverse streak orientations and magnitudes, and at the same time have complex imaging conditions such as night scenes. By using our proposed deraining algorithm MSPFN as well as other top-performing deraining methods , the restoration procedures are directly implemented on these three datasets to produce the rain-free images. And then we apply the public available pre-trained models of YOLOv3 (for detection), Mask R-CNN (for instance segmentation), and RefineNet (for semantic segmentation) to perform the the downstream tasks. Qualitative results, including the deraining performance as well as the precision of the subsequent detection and segmentation tasks, are tabulated in Table 6. In addition, visual comparisons are shown in Fig. 7.
It is obvious that rain streaks can greatly degrade the detection accuracy and segmentation precision, night scenarios in particular, i.e. by missing targets and producing low detection or segmentation confidence (mean pixel accuracy (mPA) and mean Intersection of Union (mIoU)). In addition, the detection precision of the produced rain-free images by MSPFN shows a notable improvement over that of original rain inputs by nearly 10%, and MSPFN achieves the best results of mPA as well as mIoU for semantic segmentation task on BDD150. When compared with other top-performing deraining models, the rain-free images generated by MSPFN show more credible contents with more details, which effectively promote the detection and segmentation performance. Moreover, we also evaluate our lightweight deraining model MSPFN∗ with lighter depth (, ) and width (with all filter channels of 32) since computation efficiency is crucial for mobile devices and applications require real-time throughput such as autonomous driving. MSPFN∗ still achieves competitive performance compared with other models while it’s a half more efficient in terms of inference time.
Conclusion
In this paper, we propose a novel multi-scale progressive fusion network (MSPFN) to exploit the multi-scale rain information to cooperatively represent rain streaks based on the pyramid framework. To achieve this goal, we design several basic modules (CFM, FFM and RM) along with our proposed multi-scale progressive fusion mechanism to explore the inherent correlations of the similar rain patterns among multi-scale rain streaks. Consequently, our predicted rain streak distribution is potentially more correct due to the collaborative representation of rain streaks across different scales. Experimental results on several synthetic deraining datasets and real-world scenarios, as well as several downstream vision tasks (i.e. object detection and segmentation) have shown great superiority of our proposed MSPFN algorithm over other top-performing methods.
Acknowledgement
This work is supported by National Key R&D Project (2016YFE0202300) and National Natural Science Foundation of China (U1903214, 61671332, U1736206, 41771452, 41771454, 61971165), and Hubei Province Technological Innovation Major Project (2019AAA049, 2018CFA024).