A Model-driven Deep Neural Network for Single Image Rain Removal

Introduction

Images taken under various rain conditions often suffer from unfavorable visibility, and always severely affect the performance of outdoor computer vision tasks, such as objection tracking , video surveillance , and pedestrian detection . Hence, removing rain streaks from rainy images is an important pre-processing task and has drawn much research attention in the recent years .

In the past years, various methods have been proposed for single image rain removal task. Many researchers made focus on exploring physical properties of rain layer and background layer, and introduced various prior structures to regularize and separate them. Along this research line, the representative methods include layer priors with Gaussian mixture model (GMM) , discriminative sparse coding (DSC) , and joint convolutional analysis and synthesis sparse representation (JCAS) . Especially, inspired by the fact that rain streaks repeatedly appear at different locations over a rainy image with similar local patterns like shape, thickness, and direction, very recently researchers represented this configuration of rain layer by the convolutional dictionary learning model . Such a representation finely delivers this prior knowledge by imposing rain kernels (conveying repetitive local patterns) on sparse rain maps, as intuitively depicted in Fig. 1 (a). These methods thus achieved state-of-the-art (SOTA) performance when the background can also be well represented, e.g., by low-rank prior in surveillance video sequences .

Albeit effective in certain applications, the rationality of these techniques depends on the subjective prior assumptions imposed on the unknown background and rain layers to be recovered. In real scenarios, however, such learning regimes could not always adapt to different rainy images with complex, diverse, and variant structures collected from different resources. Besides, these methods generally need time-consuming iterative computations, often with efficiency issue in real applications.

Driven by the significant success of deep learning (DL) in low level vision, recent years have also witnessed the rapid progress of deep convolutional neural networks (CNN) for single image rain removal . The current DL-based derainers mainly focus on designing network modules, and then train network parameters based on abundant rainy/clean image pairs to extract the background layer. Typical deraining network structures include deep detail network (DDN) , recurrent squeeze-and-excitation context aggregation module (RESCAN) , progressive image deraining network (PReNet) , spatial attentive unit (SPANet) , and many others.

These DL strategies, however, also possess evident deficiencies. The most significant one is their weak interpretability. Network structures are always complicated and diverse, making it difficult to analyze the role of different modules and understand the underlying insights of their mechanism. Besides, most of them treat CNN as an encapsulated end-to-end mapping module without deepening into the rationality, and neglect the intrinsic prior knowledge of rain streaks such as sparsity and nonlocal similarity. This makes this methodology easily trapped into the overfitting-to-training-sample issue.

To alleviate the aforementioned issues, this paper designs an interpretable deep network, which sufficiently considers the characteristics of rain streaks and attempts to combine the advantages of the conventional model-driven prior-based and current data-driven DL-based methodologies. Specifically, our contributions are mainly three-fold:

Firstly, we propose a concise rain convolutional dictionary (RCD) model for single image by exploiting the intrinsic convolutional dictionary learning mechanism to encode rain shapes, and specifically adopt the proximal gradient technique to design an optimization algorithm for solving it. Different from traditional solvers for the RCD model containing complex operations (e.g., Fourier transformation), the algorithm only contains simple computations (see Fig. 1 (b)) easy to be implemented by general network modules. This facilitates our algorithm capable of being easily unfolded into a deep network architecture.

Secondly, by unfolding the algorithm, we design a new deep network architecture for image deraining, called RCDNet. The specificity of this network lies on its exact step-by-step corresponding relationship between its modules and the algorithm operators, and thus successively possesses the interpretability of all its modules as that of all steps in the algorithm. Specifically, as shown in Fig. 1 (b) and (c), each iteration of the algorithm contains two sub-steps, respectively updating the rain map (convoluted by the learned rain kernels) and background layer, and each stage of the RCDNet also contains two sub-networks (M-net and B-net). Each output of the intermediate layer in the network is thus with clear interpretation, which greatly facilitates a deeper analysis on what happens inside the network during training, and a comprehensive understanding why the network works or not (as the analysis presented in Sec. 5.2).

Thirdly, comprehensive experimental results substantiate the superiority of the RCDNet beyond SOTA conventional prior-based and current DL-based methods both quantitatively and visually. Especially, attributed to its well interpretability, not only the underlying rationality and insights of the network can be intuitively understood through visualizing the amelioration process (like the gradually rectified background and rain maps) over all network layers by general users, but also the network can yield generally useful rain kernels for expressing rain shapes and proximal operators for delivering the prior knowledge of background and rain maps for a rainy image, facilitating their general availability to more real-world rainy images.

The paper is organized as follows. Sec. 2 reviews the related rain removal work. Sec. 3 presents the RCD model for rain removal as well as the algorithm designed for solving it. Then Sec. 4 introduces the unfolding deep network for the algorithm. The experimental results are demonstrated in Section 5 and the paper is finally concluded.

Related work

In this section, we give a brief review on the most related work on rain removal for images. Depending on the input data, the existing algorithms can be categorized into two groups: video based and single image based ones.

Garg and Nayar first tried to analyze the visual effects of raindrops on imaging systems, and utilized a space-time correlation model to capture the dynamics of raindrops and a physics-based motion blur model to illustrate the photometry of rain. For better visual quality, they further proposed to increase the exposure time or reduce the depth of field of a camera . Later, both temporal and chromatic properties of rain were considered and then background layer was extracted from rainy video by utilizing different strategies such as K-means clustering , Kalman filter , and GMM . Besides, a spatio-temporal frequency based raindrop detection method was provided in .

In recent years, researchers introduced more intrinsic characteristics of rainy video to the task, e.g., similarity and repeatability of rain streaks , low-rankness among multi-frames , and sparsity and smoothness of rain streaks . To handle heavy rain and dynamic scenes, a matrix decomposition based video deraining algorithm was presented in . Afterwards, rain streaks were encoded as a patch based GMM to adapt a wider range of rain variations . More characteristics of rain streaks in a rainy video were explored including repetitive local patterns and multi-scale configurations and they were described as multiscale convolutional sparse coding model . More recently, there are some DL-based methods proposed for this task. Chen et al. presented a CNN architecture and utilized superpixel to handle torrential rain fall with opaque streak occlusions. To further improve visual quality, Liu et al. designed a joint recurrent rain removal and reconstruction network that integrated rain degradation classification, rain removal, and background details reconstruction. To handle dynamic video contexts, they further developed a dynamic routing residue recurrent network . Though these methods work well for videos, they cannot directly perform in single image cases due to the lack of temporal knowledge.

2 Single image deraining methods

Compared with video deraining task under a sequence of images, rain removal from a single image is much more challenging. The early attempts utilized the model-driven strategies by decomposing a single rainy image into low frequency part (LFP) and high frequency part (HFP) and then specifically extracted rain layer from the HFP based on various processing such as guided filter and nonlocal means filtering . Later, researchers made more focus on exploring the prior knowledge of rain and rain-free layers of a rainy image, and designing proper regularizer to extract and separate them . E.g., considered the specific sparsity characteristics of rain-free and rain parts and expressed them as the joint analysis and synthesis sparse representation models, respectively. used a similar manner to deliver local repetitive patterns of rain streaks across the image as the RCD model. Albeit achieving good performance on certain scenarios, these prior-based methods rely on the subjective prior assumptions, while could not always generally work well for practical complicated and highly diverse rain shapes in real rainy images collected from different resources.

Recently, a number of DL-based single image rain streak removal methods were proposed through constructing diverse network modules . To handle heavy rain, Yang et al. developed a multi-stage joint rain detection and estimation network for single image (JORDER_E). Very recently, Ren et al. designed a PReNet that repeatedly unfolded several Resblocks and a LSTM layer. Wang et al. presented an attention unit based SPANet for removing rain in a local-to-global manner. Through using abundant rainy/clean image pairs to train the deep model, these methods achieve favorable visual quality and SOTA quantitative measures of derained results. Most of these methods, however, just utilize network modules assembled with some off-the-shelf components in current DL toolkits to directly learn background layer in an end-to-end way, and largely ignore the intrinsic prior structures inside the rain streaks. This makes them lack of evident interpretability in their network architectures and still have room for further performance enhancement.

At present, there is a new type of single image derainers that try to combine prior and DL methodologies. For example, Mu et al. utilized CNN to implicitly learn prior knowledge for background and rain streaks, and formulated them into traditional bi-layer optimization iterations. Wei et al. provided a semi-supervised rain removal method (SIRR) that described rain layer prior as a general GMM and jointly trained the backbone–DDN. Albeit obtaining initial success, they still use CNN architectures as their main modules to construct the network, which is thus still lack of sufficient interpretability.

RCD model for single image deraining

where $\mathcal{B}$ and $\mathcal{R}$ represent the background and rain layers of the image, respectively. Then, the aim of most of DL-based deraining methods is to estimate the mapping function (expressed by a deep network) from $\mathcal{O}$ to $\mathcal{B}$ (or $\mathcal{R}$).

Instead of heuristically constructing a complex deep network architecture, we first consider the problem under the conventional prior-based methodology through exploiting the prior knowledge for representing rain streaks . Specifically, as shown in Fig. 1 (a), by adopting the RCD mechanism, the rain layer can be modeled as:

It should be noted that the rain kernels actually can be viewed a set of convolutional dictionary for representing repetitive and similar local patterns underlying rain streaks, and a small number of rain kernels can finely represent wide range of rain shapesWe simply set $N=32$ for all our experiments.. They are common knowledge for representing different rain types across all rainy images, and thus could be learned from abundant training data by virtue of the strong learning capability of end-to-end training manner of deep learning (see more details in Sec. 4). Unlike rain kernels, the rain maps must vary with the input rainy image as the locations of rain streaks are totally random. Therefore, for predicting the clean image from a testing input rainy one, the key issue is to output $\bm{M}_{n}$s and $\mathcal{B}$ from $\mathcal{O}$ with the rain kernels $\mathcal{C}_{n}$s fixed, and the corresponding optimization problem is:

2 Optimization algorithm

Since we want to build a possibly perfect step-by-step corresponding deep unfolding network architecture for solving the problem (4), it is critical to build an algorithm which contains only simple computations easy to be transformed to network modules. The traditional solvers for RCD-based model usually contain certain complicated operations, e.g., the Fourier transform and inverse Fourier transform , which are hard to accomplish such exact transformation from algorithm to network structure. We thus prefer to build a new algorithm for solving the problem through alternately updating $\mathcal{M}$ and $\mathcal{B}$ by proximal gradient method . In this manner, only simple computations can be involved. The details are as follows:

Updating $\mathcal{M}$: The rain maps $\mathcal{M}$ can be updated by solving the quadratic approximation of the problem (4) as:

where $\mathcal{M}^{(s-1)}$ is the updating result of the last iteration, $\eta_{1}$ is the stepsize parameter, and $f\left(\mathcal{M}^{(s-1)}\right)\!=\!\left\|\mathcal{O}\!-\!\mathcal{B}^{(s-1)}\!-\!\sum_{n=1}^{N}{\mathcal{C}}_{n}\!\otimes\!\bm{M}^{(s-1)}_{n}\right\|_{F}^{2}$. Corresponding to general regularization terms , the solution of Eq. (5) is:

where $\mbox{prox}_{\alpha\eta_{1}}(\cdot)$ is the proximal operator dependent on the regularization term $g_{1}(\cdot)$ with respect to $\mathcal{M}$. Instead of choosing a fixed regularizer in the model, the form of the proximal operator can be automatically learned from training data. More details will be presented in the next section.

Updating $\mathcal{B}$: Similarly, the quadratic approximation of the problem (4) with respect to $\mathcal{B}$ is:

where $\nabla h\left(\!\mathcal{B}^{(s-1)}\!\right)=\sum_{n=1}^{N}\mathcal{C}_{n}\otimes\bm{M}_{n}^{(s)}+\mathcal{B}^{(s-1)}-\mathcal{O},$ and it is easy to deduce that the final updating rule for $\mathcal{B}$ is3:

where $\mbox{prox}_{\beta\eta_{2}}(\cdot)$ is the proximal operator correlated to the regularization term $g_{2}(\cdot)$ with respect to $\mathcal{B}$.

Based on this iterative algorithm, we can then construct our deep unfolding network as follows.

The rain convolutional dictionary network

Inspired by the recently raised deep unfolding techniques in various tasks such as deconvolution , compressed sensing , and dehazing , we build a network structure for single image rain removal task by unfolding each iterative steps of the aforementioned algorithm as the corresponding network module. We especially focus on making all network modules one-to-one corresponding to the algorithm implementation operators, for better interpretability.

As shown in Fig. 2 (a), the proposed network consists of $S$ stages, corresponding to $S$ iterations of the algorithm for solving (4). Each stage achieves the sequential updates of $\mathcal{M}$ and $\mathcal{B}$ by M-net and B-net. As displayed in Fig. 2 (b), exactly corresponding to each iteration of the algorithm, in each stage of the network, M-net takes the observed rainy image $\mathcal{O}$ and the previous outputs $\mathcal{B}^{(s-1)}$ and $\mathcal{M}^{(s-1)}$ as inputs, and outputs an updated $\mathcal{M}^{(s)}$, and then B-net takes $\mathcal{O}$ and $\mathcal{M}^{(s)}$ as inputs, and outputs an updated $\mathcal{B}^{(s)}$.

The key issue of unrolling the algorithm here is how to represent the two proximal operators involved in (8) and (10) while other operations can be naturally performed with generally used operators in normal networks . In this work, we simply choose a ResNet to construct the two proximal operators as many other works did . Then, we can separately decompose the updating rules for $\mathcal{M}$ as (8) and $\mathcal{B}$ as (10) into sub-steps and achieve the following procedures for the $s^{\text{th}}$ stage of the RCDNet:

where $\text{proxNet}_{\theta_{m}^{(s)}}(\cdot)$ and $\text{proxNet}_{\theta_{b}^{(s)}}(\cdot)$ are two ResNets consisting of several Resblocks with the parameters ${\theta_{m}^{(s)}}$ and ${\theta_{b}^{(s)}}$ at the $s^{\text{th}}$ stage, respectively.

We can then design the network architecture, as shown in Fig. 2, by transforming the operators in (11) and (12) step-by-step. All the parameters involved can be automatically learned from training data in an end-to-end manner, including $\{\theta_{m}^{(s)},\theta_{b}^{(s)}\}_{s=1}^{S}$, rain kernels $\mathcal{C}$, $\eta_{1}$, and $\eta_{2}$.

It should be indicated that both of the two sub-networks are very interpretable. As shown in Fig. 2 (b), the M-net accomplishes the extraction of residual information $\mathcal{E}^{(s)}$ of rain maps. Specifically, $\widehat{\mathcal{R}}^{(s)}$ is the rain layer estimated with the previous background $\mathcal{B}^{(s-1)}$, and $\widetilde{\mathcal{R}}^{(s)}$ is the rain layer achieved by the generative model (2) with the estimated $\mathcal{M}^{(s-1)}$. Then the M-net calculates the residual information between the two rain layers obtained in this two ways, and extracts the residual information $\mathcal{E}^{(s)}$ of rain maps with the transposed convolution of rain kernels to update the rain map. Next, the B-net recovers the background $\widehat{\mathcal{B}}^{(s)}$ estimated with current rain kernel and rain maps $\mathcal{M}^{(s)}$, and fuses this estimated $\widehat{\mathcal{B}}^{(s)}$ with the previously estimated $\mathcal{B}^{(s-1)}$ by weighted parameters $\eta_{2}$ and ($1-\eta_{2}$) to get the updated background ${\mathcal{B}}^{(s)}$. Here, we set $\mathcal{M}^{(0)}$ as 0 and initialize $\mathcal{B}^{(0)}$ by a convolutional operator on $\mathcal{O}$More network design details are described in supplemental file..

Remark: From Fig. 2, the input tensor of $\text{proxNet}_{\theta_{b}^{(s)}}(\cdot)$ has the same size ${H\times W\times 3}$ as the to-be-estimated $\mathcal{B}$. Evidently, this is not beneficial for learning $\mathcal{B}$ since most of the previous updating information would be compressed due to few channels. To better keep and deliver image features, we simply expand the input tensor at the $3^{\text{rd}}$ mode for more channels in experiments (see more in supplemental file).

2 Network training

Training loss. For simplicity, we adopt the mean square error (MSE) for the learned background and rain layer at every stage as the training objective function:

where $\mathcal{B}^{(s)}$ and $\mathcal{R}^{(s)}$ separately denote the derained result and extracted rain layer as expressed in (12) at the $s^{\text{th}}$ stage ($s=0,1,\cdots,S$). $\lambda_{s}$ and $\gamma_{s}$ are tradeoff parametersIn all experiments, we simply set $\lambda_{S}=\gamma_{S}=1$ to make the outputs at the final stage play a dominant role, and other parameters as 0.1 to help find the correct parameter in each stage. More parameter settings are discussed in supplementary material..

Implement details. We implement our network based on a NVIDIA GeForce GTX 1080Ti GPU. We adopt the Adam optimizer with the batch size of 16 and the patch size of 64$\times$64. The initial learning rate is $1\times$$10^{-3}$ and divided by 5 every 25 epochs. The total epoch is 100.

Experimental results

We first conduct ablation study and model visualization to verify the underlying mechanism of the proposed network, and then present experiments on synthesized benchmark datasets and real datasets for performance evaluation.

Dataset and performance metrics. In this section, we use Rain100L to perform all the ablation studies. The synthesized dataset consists of 200 rainy/clean image pairs for training and 100 pairs for testing . Two performance metrics are employed, including peak-signal-to-noise ratio (PSNR) and structure similarity (SSIM) . Note that as the human visual system is sensitive to the Y channel of a color image in YCbCr space, we compute PSNR and SSIM based on this luminance channel.

Table 1 reports the effect of stage number $S$ on deraining performance of our network. Here, $S=0$ means that the initialization $\mathcal{B}^{(0)}$ is directly regraded as the recovery result. Taking $S=0$ as a baseline, it is seen that only with 2 stages, our method achieves significant rain removal performance, which validates the essential role of the proposed M-net and B-net. We also observe that when $S=20$, its deraining performance is slightly lower than that of $S=17$ since larger $S$ would make gradient propagation more difficult. Based on such observation, we easily set $S$ as 17 throughout all our experiments. More ablation results and discussions are listed in supplementary material.

2 Model verification

We then show how the interpretability of this RCDNet facilitates an easy analysis for the working mechanism inside the network modules.

Fig. 3 presents the extracted background layer $\mathcal{B}^{(s)}$ ($1^{\text{st}}$ row), $\widehat{\mathcal{B}}^{(s)}$($2^{\text{nd}}$ row) that represents the role of M-net in helping restore clean background, and rain layer ${\mathcal{R}}^{(s)}$ ($3^{\text{rd}}$ row) at different stages. We can find that with the increase of $s$, ${\mathcal{R}}^{(s)}$ covers more rain streaks and fewer image details, and $\widehat{\mathcal{B}}^{(s)}$ and $\mathcal{B}^{(s)}$ are also gradually ameliorated. These should be attributed to the proper guidance of the RCD prior for rain streaks and the mutual promotion of M-net and B-net that enables the RCDNet to be evolved to a right direction.

Fig. 4 presents the learned rain kernels and the rain maps for the input $\mathcal{O}$ in Fig. 3. Clearly, the RCDNet finely extracts proper rain layers explicitly based on the RCD model. This not only verifies the reasonability of our method but also manifests the peculiarity of our proposal. On one hand, we utilize a M-net to learn sparse rain maps instead of directly learning rain streaks that makes learning process easier. On the other hand, we exploit training data to automatically learn rain kernels representing general repetitive local patterns of rain with diverse shapes. This facilitates their general availability to more real-world rainy images.

3 Experiments on synthetic data

Comparison methods and datasets. We then compare our network with the current SOTA single image derainers, including model-based DSC , GMM , and JCAS ; DL-based Clear , DDN , RESCAN , PReNet , SPANet , JORDER_E , and SIRR The code/project links for these comparison methods are listed in supplementary material., based on four benchmark datasets, including Rain100L, Rain100H , Rain1400 , and Rain12 .

Fig. 5 illustrates the deraining performance of all competing methods on a rainy image from Rain100L. As shown, the deraining result of RCDNet is better than that of other methods in sufficiently removing the rain streaks and finely recovering the image textures. Moreover, the rain layer extracted by RCDNet contains fewer unexpected background details as compared with other competing methods. Our RCNet thus achieves the best PSNR and SSIM.

Table 2 reports the quantitative results of all competing methods. It is seen that our RCDNet attains best deraining performance among all methods on each dataset. This substantiates the flexibility and generality of our method, in diverse rain types contained in these datasets.

4 Experiments on real data

We then analyze the performance of all methods on two real datasets from : the first one (called SPA-Data) contains 638492 rainy/clean image pairs for training and 1000 testing ones, and the second one (called Internet-Data) includes 147 rainy images without groundtruth.

Table 3 and Fig. 6 compare the derained results on SPA-Data of all competing methods visually and quantitatively. It is easy to see that even for such complex rain patterns, the proposed RCDNet still achieves an evident superior performance than other methods. Especially, similar to its superiority in synthetic experiments, it is also observed that our method better removes the rain streaks and recovers image details than other competing ones.

Further, we select two real hard samples with various rain densities to evaluate the generalization ability of all competing methods. From Fig. 7, we can find that traditional model-based methods tend to leave obvious rain streaks. Although DL-based comparison methods remove apparent rain streaks, they still leave distinct rain marks or blur some image textures. Comparatively, our RCDNet better preserves background details as well as removes more rain streaks. This shows its good generalization capability to unseen complex rain types.

Conclusion

In this paper, we have explored the intrinsic prior structure of rain streaks that can be explicitly expressed as convolutional dictionary learning model, and proposed a novel interpretable network architecture for single image deraining. Each module in the network can one-to-one correspond to the implementation operators of the algorithm designed for solving the model, and thus the network is almost “white-box” with easily visualized interpretation for all its module elements. Comprehensive experiments implemented on synthetic and real rainy images validate that such interpretability brings a good effect of the proposed network, and especially facilitates the analysis for how it happens in the network and why it works in testing prediction process. The extracted elements through the end-to-end learning by the network, like the rain kernels, are also potentially useful for the related tasks on rainy images.

Acknowledgment. This research was supported by the China NSFC projects under contract 11690011, 61721002, U1811461 and MoE-CMCC “Artifical Intelligence” Project with No. MCM20190701

References