DANet: Image Deraining via Dynamic Association Learning

1 Introduction

Images taken outdoors are more susceptible to rain weather conditions which decrease the visual quality of the captured images. Both rain streaks and rain accumulation can cause severe occlusion on the background scene and significantly degrade the contrast and visibility. These negative visual effects result in the degradation of many high-level computer vision tasks like image understanding, object detection [Xu et al., 2021b] and identification [Xu et al., 2021a]. Therefore, there is a pressing need to develop algorithms that remove perturbation while recovering the clean background textures.

∗Corresponding Author

Past decades have witnessed the significant progress of image deraining technologies. The early model-based methods [Garg and Nayar, 2005] rely more on statistical analyses of rain streaks and background scenes, and enforce handcrafted priors (e.g., sparsity and nonlocal means filtering) on both rain and background. Since the restrictive assumptions are introduced for specific scenes or rain patterns, these methods [Kang et al., 2012] become still biased and less generalizable on complex or real scenarios.

More recently, data-driven methods have overtaken modelbased methods on both deraining performance and popularity. These methods promote further progress by sophisticated architectures [Jiang et al., 2021] and training practices [Jiang et al., 2020a] for effective and efficient models. However, similar to image denoising, most of the prevalent models regard the image deraining as a simple rain streaks removal problem, i.e., they produce deraining results by subtracting the predicted rain residual from rainy input based on the simple additive model. Consequently, the background contents suffer from incomplete and undesirable restoration, such as over-smoothing, halo artifacts, and color distortion, due to ignoring the additional degradation side effects. Recurrent- and cascaded-based methods [Ren et al., 2019; Deng et al., 2020] have made further steps on fidelity and textural plausibility by approaching both rain streaks removal and detail recovery. As shown in Figure 1, the cascaded (DRDNet) or recurrent learning (PreNet) method could bring considerable improvements, but is still far from generating desirable results under heavy rain conditions since the rain estimation and detail recovery are isolated. In addition, they perform inference in the whole resolution space using several cascaded stages or sub-networks, which limits their representation efficiency and interpretability since the rain streaks in rainy images are commonly sparse and have the canonical samples.

Indeed, rain streaks removal and texture restoration are highly correlated, where the predicted rain distribution maps essentially reveal the degradation location and intensity, which are helpful for the accurate background restoration. This motivates us to study the following key question. Question One: Whether the predicted rain distribution can provide degradation priors to help background texture recovery?

Inspired by the dynamic network, it can adapt structures or parameters to their input during inference, leading to notable advantages in terms of accuracy, computational efficiency and adaptiveness [Han et al., 2021]. Because of the sparse distribution of rain streaks in rainy scenes, the feature-selection dynamic strategy provides the alternative solution for accurate image deraining and background restoration. Specifically, it allows the network to focus on the informative region through the pixel-wise dynamic weights (rain degradation mask), thus enjoying favorable properties those are absent in aforementioned deraining models. This way, it is reasonable to employ the dynamic learning to predict the degradation mask from the rain input to guide the accurate texture restoration. We thus construct a novel dynamic associated network (DANet) for single image deraining. DANet is a two-stage framework, whose key idea is to cast the image deraining task as a composition of rain streaks removal, texture (detail and contrast) recovery, and their association learning. More specifically, the dynamic degradation prior is learned directly from the rainy input, and then produces the degradation mask to help extract the complementary information for accurate texture restoration. In detail, a deraining sub-network (DSN) and a reconstruction sub-network (RSN) are respectively designed to approach these two sub-problems, respectively. Meanwhile, an associated learning module (ALM) is used to associate these two sub-tasks. It provides the dynamic prediction in terms of the location and degradation intensity from rainy inputs, which is more helpful for the accurate texture restoration. To achieve the optimal approximation, we introduce joint constraints for enhancing the compatibility of deraining model with background reconstruction, automatically learned from the training data.

In addition, we observe that the reconstructed rainy image via bilinear interpolation from the sampled rainy image has the similar statistical distribution to the original one, shown in Figure 2. This raises another key question. Question Two: Whether the rain distribution can be inferred from the sampling space? Simultaneously, the background details follow this observation.

This inspires us to learn the prediction of rain streak distribution and degradation mask at the sampling space, and then exploit the super-resolution to reconstruct the spatial details. There are two substantial superiorities: 1) predicting and characterizing rain distribution in the sub-space from the sampled low-resolution rainy image can greatly simplify the first learning task; 2) all the feature representations, including DSN and RSN are processed at the low-dimension space, greatly saving the computational cost. In particular, besides introducing the association learning and super-resolution technologies to alleviate the information loss due to spatial sampling, compared to the efficiency proceeds, we have verified that the performance decline is acceptable. Overall, the main contributions are summarized as follows.

1. We propose a novel dynamic association network (DANet) for image deraining. It consists of the rain streak removal, background restoration and their association learning. To the best of our knowledge, we are the first to approach image deraining from a new perspective of joint rain streaks removal and texture reconstruction via dynamic association learning and optimization.
2. We design a novel associated learning module (ALM) to associate rain streaks removal and texture recovery tasks via the dynamic degradation prior learning. It significantly alleviates the learning burden while promoting recovery accuracy. We propose a selective fusion block (SFB) for effective multi-scale residue fusion.
3. Comprehensive experiments on image deraining and detection task have verified that our proposed DANet method is more suitable for preserving the textural details while removing rain streaks.

2 ProposedMethod

This section details the architecture of our proposed dynamic association network (DANet) and its essential components.

2.1 Architecture andModel Optimization

Figure 3 outlines the framework of our proposed DANet, which contains a deraining sub-network (DSN), an associated learning module (ALM), and a reconstruction sub-network (RSN). For convenience, DSN and RSN share the same dualbranch fusion network, elaborated in Section 2.2: Dualbranch Fusion Network.

For a given rainy image IRain and its clean version IB, we first sample them with Bilinear operation to generate the corresponding sub-samples (IRain,S and IB,S). IRain,S is then input to DSN to predict the distribution of rain streaks and produce the corresponding deraining result via GDSN(·). The above procedures can be expressed as

To promote texture recovery, we propose the associated learning module (ALM) based on dynamic learning, shown in Figure 3. ALM is designed to associate rain streaks removal and texture recovery tasks. Specifically, it takes the predicted derained image I ∗B,S , rain residue I ∗R,S at sub-space and the original rainy image IRain as inputs, and first learns a global rain-distribution mask from I ∗R,S . The mask provides global degradation priors (location and density) and allows the network to borrow background texture information fti from IRain. A novel selection fusion block (SFB), described in Sec. 2.3 and Supplementary, is then used to distill the features for enhancement. These procedures of ALM are expressed as

In Equation (2), FR(·) denotes the feature embedding function, including an initial layer and a strided convolution layer to learn the embedding representation of IRain. FM(·) refers to the mask function, which is composed of an initial layer and a bottleneck unit to learn the rain-distribution matrices. A Sigmoid operator is followed to normalize them into [0, 1.0]. Through the positional guidance and weights, we can borrow the rain-free texture information fti from the embedding representation of IRain through the point-wise multiplication ◦. The visualization of the mask and fti are provide in Figure 4. FB(·) is the embedding function to generate the initial representation of I ∗B,S . FSFB(·) denotes the fusion function in SFB. Following that, RSN takes fALM as input to achieve the texture reconstruction as follows:

where GRSN(·) denotes the super-resolving function of RSN, and FUP (·) is the Bilinear interpolation.

To train these two sub-networks jointly, we introduce the association optimization, including the image loss (Charbonnier penalty loss [Lai et al., 2017]) and a structural similarity (SSIM) loss, which supervise networks to achieve image and structural fidelity restoration simultaneously. The loss functions are given by

where α and λ are used to balance the loss components, and experimentally set as -0.15 and 1, respectively. The penalty coefficient ε is set to 10-3 .

2.2 Dual-branch Fusion Network

The backbone of DSN and RSN is a deep dual-branch fusion network, involving an original resolution branch (ORB) and an encoder-decoder branch. The former is used to preserve spatial structure from the input to the output image. The latter is designed to infer locally-enriched textures.

2.3 Selective Fusion Block

Considering the feature redundancy and knowledge discrepancy among residual blocks and encoding stages, we introduce a novel selective fusion block (SFB) where the lowlevel contextualized features of earlier stages help consolidate high-level features of the later stages (or scales). Specifically, we incorporate depth-wise separable convolutions and the channel attention layer into SFB to discriminatively aggregate multi-scale features in spatial and channel dimensions. More analyses and discussions to verify its effectiveness are presented in Section 3.2: Ablation Study.

3 Experiments

To validate our proposed DANet, we conduct extensive experiments on synthetic and real-world rainy datasets, and compare DANet with eleven image deraining methods. These methods include MPRNet [Zamir et al., 2021], SWAL [Huang et al., 2021], DRDNet [Deng et al., 2020], MSPFN [Jiang et al., 2020b], IADN [Jiang et al., 2020a], PreNet [Ren et al., 2019], UMRL [Yasarla and Patel, 2019], DIDMDN [Zhang and Patel, 2018], RESCAN [Li et al., 2018], LPNet [Fu et al., 2020] and DDC [Li et al., 2019b]. Five evaluation metrics, such as Peak Signal to Noise Ratio (PSNR), Structural Similarity (SSIM), Feature Similarity (FSIM), Naturalness Image Quality Evaluator (NIQE) and Spatial-Spectral Entropy-based Quality (SSEQ) [Liu et al., 2014], are employed for comparison.

3.1 Implementation Details

Data Collection. Since there exits the discrepancy in training samples for all comparison methods, following [Jiang et al., 2020b], we use 13, 700 clean/rain image pairs from [Zhang et al., 2020; Fu et al., 2017] for training all comparison methods with their publicly released codes by tuning the optimal settings for a fair comparison. For testing, four synthetic benchmarks (Test100 [Zhang et al., 2020], Test1200 [Zhang and Patel, 2018], R100H and R100L [Yang et al., 2017]) and three real-world datasets (Rain in Driving (RID), Rain in Surveillance (RIS) [Li et al., 2019a] and Real127 [Zhang and Patel, 2018]) are considered for evaluation.

Experimental Setup. In our baseline, the number of RCAB is empirically set to 2 for each stage in the encoder-decoder branch and 5 for the original resolution branch with filter numbers of 48. The training images are coarsely cropped into small patches with a fixed size of 128 × 128 pixels to obtain the training samples. We use Adam optimizer with the learning rate (4 × 10-4 with the decay rate of 0.8 at every 80 epochs till 500 epochs) and batch size (16) to train our DANet on a single NVIDIA Titan Xp GPU.

3.2 Ablation Study

Validation on Basic Components. We conduct ablation studies to validate the contributions of individual components, including the depth-wise separable convolutions (DSC), selective fusion block (SFB) and associated learning module (ALM) to the final deraining performance. For simplicity, we denote our final model as DANet and devise the baseline model w/o all by removing all these components above. Quantitative results in terms of deraining performance and inference efficiency on the Test1200 dataset are presented in Table 1, revealing that the complete deraining model DANet achieves significant improvements over its incomplete versions. Compared to w/o ALM model (removing ALM from DANet), DANet achieves 2.45dB performance gain since the association learning in ALM can help the network to fully exploit the background texture from the rainy image under the guidance of predicted rain distribution. In addition, disentangling the image deraining task into rain streaks removal and texture reconstruction at the low-dimension space exhibits considerable superiority in terms of efficiency (26.2% and 70.6% more efficient in inference time and computational cost, respectively) and restoration quality (referring to the results of DANet and DANet∗ models (completing the deraining and texture recovery at the original resolution space [Deng et al., 2020])). Moreover, using the depth-wise separable convolutions allows increasing the channel depth with approximately the same parameters, thus enhancing the representation capacity.

We also conduct ablation studies to validate the dualbranch fusion, involving an original resolution branch (ORB) and a U-shaped encoder-decoder branch (EDB). Based on DANet, we devise two comparison models (w/o ORB and w/o EDB) by removing these two branches in turn. Quantitative results are presented in Table 1. Removing ORB may greatly weaken the representation capability on the spatial structure, leading to obvious performance decline (1.15dB in PSNR) (referring to the results of DANet and w/o ORB models). Moreover, EDB allows the network to aggregate multiscale textures, which are crucial for fine detail recovery.

3.3 Comparison with State-of-the-arts

Synthesized Data. Quantitative results on Test1200 and Test100 are provided in Table 2. As expected, our highaccuracy DANet♠ model (by setting the number of RCAB in ORB to 10 with the channel depth of 64) achieves the best and second-best evaluated scores on Test1200 and Test100. Visual comparisons in Figure 5 show that high-accuracy methods, such as PreNet, UMRL and MSPFN, can effectively eliminate the rain layer and thus bring an improvement in visibility. But they fail to generate visual appealing results due to the inaccurate color correction. Likewise, DRDNet focuses on the detail recovery, but shows undesired deraining performance. Besides recovering cleaner and more credible image textures, our DANet produces results with better contrast and less color distortion. Please refer to the "skiing" and "airplane" scenarios. These results further demonstrate the effectiveness of our proposed dynamic association learning.

The results on R100H/R100L datasets are provided in Table 2, indicating that most of the deraining models obtain impressive performance on light rain cases with high consistency. However, only our DANet and MPRNet still perform favorably on heavy rain conditions, exhibiting great superiority over other competing methods in terms of PSNR. Note that MPRNet achieves limited superiority (average 0.29dB on 4 datasets) over our DANet♠, but requires 89.9%, 23.6% and 332.2%more inference time, parameters and computational cost. For more convincing evidence, we also provide additional visual comparisons in Figure 5. Our DANet model can remove more rain streaks and achieve better image fidelity than the competing methods. Especially for the heavy rain condition, other deraining methods fail to eliminate the rain streaks, and introduce considerable artifacts and unnatural color appearance. (please refer to the first two scenarios in Figure 5). We speculate that these visible improvements on restoration quality may benefit from the elaborate design of the framework, including the novel deraining scheme - rain streaks removal and texture reconstruction as well as the dynamic association learning in ALM. These strategies are integrated into a unified framework, allowing the network to fully exploit the contextual information for feature representation and complete the deraining task more easily.

Real-world Data. We further conduct experiments on three real-world datasets: Real127 [Zhang and Patel, 2018], Rain in Driving (RID), and Rain in Surveillance (RIS) [Li et al., 2019a]. Quantitative results of NIQE and SSEQ [Liu et al., 2014] are listed in Table 3, where smaller NIQE and SSEQ scores indicate better perceptual quality and clearer contents. Again, our proposed DANet is highly competitive, achieving the lowest average values on the RID dataset and the best average score of SSEQ on the Real127 dataset. We visualize the deraining results in Figure 6, showing that DANet produces rain-free images with cleaner and more credible contents, whereas the competing methods fail to remove rain streaks. These evidences indicate that our DANet model performs well in separating rain information while preserving textural details and image naturalness, even under real rain conditions.

3.4 Impact on DownstreamVision Task

Eliminating the degraded effects of rain streaks under rainy conditions while preserving credible textural details is crucial for object detection. This motivates us to investigate the effect of deraining performance on object detection accuracy based on popular object detection algorithms (e.g., YOLOv3 [Redmon and Farhadi, 2018]). Based on our DANet and several representative deraining methods, the restoration procedures are directly applied to the rainy images to generate corresponding rain-free outputs. We then apply the publicly available pre-trained models of YOLOv3 for the detection task. Table 4 shows that DANet achieves the highest PSNR scores on COCO350 and BDD350 datasets [Jiang et al., 2020b]. Meanwhile, the rain-free results generated by DANet facilitate better object detection performance than other methods. Visual comparisons on two instances in Figure 7 indicate that DANet exhibit a notable superiority in terms of image quality and detection accuracy. We attribute the considerable performance improvements of both deraining and down-stream detection tasks to our association learning between rain streaks removal and detail recovery tasks.

4 Conclusion

We rethink the image deraining as a composite task of rain streak removal, textures recovery and their association learning, and propose a dynamic associated network (DANet) for image deraining. Accordingly, a two-stage architecture and an associated learning module (ALM) are adopted in DANet to account for twin goals of rain streak removal and texture reconstruction while facilitating the learning capability. Meanwhile, the joint optimization promotes the compatibility while maintaining the model compactness. Extensive results on image deraining and joint detection task demonstrate the superiority of our DANet model over the state-of-the-arts.

Acknowledgments

This work is supported by National Natural Science Foundation of China (U1903214, 62071339, 61872277, 62072347, 62171325), Natural Science Foundation of Hubei Province (2021CFB464), Open Research Fund from Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ)(GML-KF-22-16) and Guangdong-Macao Joint Innovation Project (2021A0505080008).