Single Image Deraining with Continuous Rain Density Estimation

Under the framework of CSC, one can express the mathematical model for rain streaks:

where $\left\{\boldsymbol{z}_{i}\right\}_{i=1}^{c}\in\mathbb{R}^{c\times l_{h}\times l_{w}}$ are $c$ sparse coefficients of noisy rain streak map $\boldsymbol{r}_{\epsilon}\in\mathbb{R}^{1\times l_{h}\times l_{w}}$ w.r.t. corresponding rain kernels $\left\{\boldsymbol{f}_{i}\right\}_{i=1}^{c}\in\mathbb{R}^{c\times s\times s}$, $\boldsymbol{\epsilon}$ denotes the noise and $\otimes$ denotes the convolution operation.

To yield $\boldsymbol{z}_{i}$, a common solution is to solve the following $\ell_{1}$-minimization problem:

where $\|\cdot\|_{p}$ is $\ell_{p}$-norm and $\lambda$ is a regularization parameter balancing sparsity and fidelity.

Intrinsically, the sparse coefficients could reflect the rain density since the light rain exhibits a simpler structure and thus results in coefficients sparser than heavy rain. To take into account the rain density implicitly, we introduce a weight parameter $w_{i}$ for $\boldsymbol{z}_{i}$ [12]:

Due to the fact that convolution operation can be replaced with matrix multiplication, Eq. (4) can be reformulated as:

where $\boldsymbol{F}=[\boldsymbol{F}_{1},\boldsymbol{F}_{2},\dots,\boldsymbol{F}_{c}]$ is the concatenation of convolution matrices corresponding to filters $\boldsymbol{f}_{i}$, $\boldsymbol{z}=[\boldsymbol{z}_{1},\boldsymbol{z}_{2},\dots,\boldsymbol{z}_{c}]$, and $\boldsymbol{F}_{i}\boldsymbol{z}_{i}\equiv\boldsymbol{f}_{i}\otimes\boldsymbol{z}_{i}$. Note that $\boldsymbol{r}_{\epsilon}$ and $\boldsymbol{z}_{i}$ are ordered lexicographically as column vectors. From this view, CSC can be regarded as a special case of conventional sparse coding. Thus the solution to Eq. (5) can be obtained by following iterations [12]:

where $\boldsymbol{w}=[w_{1},w_{2},\dots,w_{c}]\in\left(0,1\right]$, $\Gamma_{\theta}(\cdot)$ denotes the element-wise soft thresholding function, defined by $\Gamma_{\theta}(\alpha)=\text{sign}(\alpha)\text{max}(|\alpha|-\theta,0)$, $\boldsymbol{I}$ is the identity matrix, and $L$ is the Lipschitz constant. Note that, since $\boldsymbol{I}$ and $\boldsymbol{F}$ are sparse convolution matrices, there must be $\boldsymbol{S}$ and $\boldsymbol{G}$ satisfying $\boldsymbol{S}\otimes\boldsymbol{z}^{(t)}=\left(\boldsymbol{I}-\frac{1}{L}\boldsymbol{F}^{T}\boldsymbol{F}\right)\boldsymbol{z}^{(t)}$ and $\boldsymbol{G}\otimes\boldsymbol{r}_{\epsilon}=\frac{1}{L}\boldsymbol{F}^{T}\boldsymbol{r}_{\epsilon}$.

Since $\Gamma_{\theta}(\cdot)$ is an element-wise operation and $\boldsymbol{w}>0$, $\Gamma_{\theta}(\cdot)$ has an important attribute:

where $\boldsymbol{\tilde{w}}=[\frac{1}{{w}_{1}},\frac{1}{{w}_{2}},\dots,\frac{1}{{w}_{c}}]$, and $\odot$ denotes the element-wise product. Based on such observation, we rewrite Eq. (6) as:

In Eq. (8), $\boldsymbol{\tilde{w}}$ is mainly used to adjust sparsity of $\boldsymbol{z}$ while $\boldsymbol{w}$ for scaling. Large $\boldsymbol{\tilde{w}}$ favors less sparse solution, and vice versa. Namely, $\boldsymbol{\tilde{w}}$ is conducive to more accurate representations of rain with various densities, resulting in better deraining results. On the other hand, $\boldsymbol{\tilde{w}}$ implies the rain density and thus could be utilized to estimate rain density.

Recall the CSC model for rain streaks Eq. (2), the basic idea behind Eq. (2) is to sparsely represent the rain streaks with learned kernels $\{\boldsymbol{f}_{i}\}_{i=1}^{c}$. Intrinsically, the light rain is supposed to have more concise representation than the heavy rain due to the fact of simpler appearance. To this end, we choose the channel attention (CA) as our learning weight block (LWB), defined by:

where $Fc_{1}$ and $Fc_{2}$ are full connection layers without biases, ReLU denotes the Rectified Linear Unit activation function, $\text{AvgPool}(\cdot):\mathbb{R}^{c\times h\times w}\rightarrow\mathbb{R}^{c\times 1\times 1}$ denotes the global average pooling function, and $\delta(\cdot)$ denotes the sigmoid function.

Benefitting from Eq. (9), for light rain, small valued weights $\tilde{\boldsymbol{w}}$ in Eq. (8) would be generated, which means Eq. (9) will penalize the sparse coefficients $\boldsymbol{z}$ much more to reduce sparsity. On the contrary, for the heavy rain with less sparse representation under the learned kernels $\{\boldsymbol{f}_{i}\}_{i=1}^{c}$, the corresponding weights $\tilde{\boldsymbol{w}}$ will be large to enhance sparsity. As a consequence, by adaptively weighting the sparse coefficients in Eq. (8) using LWB Eq. (9), the rain density can be implicitly considered, which proved in Section V contributes to better rain removal results.

On the other hand, since the rain density is closely related to the weights $\tilde{\boldsymbol{w}}$, we could estimate the rain density from $\tilde{\boldsymbol{w}}$. As shown in Fig. 4(d), we illustrate the distribution of learned weights $\tilde{\boldsymbol{w}}$ in Eq. (8), one can see that the weight cluster gradually approaches 1 with the increase of rain intensity. Therefore, we define the average $\tilde{\boldsymbol{w}}$ of different channels as rain density estimation (RDE):

the more it rains, the greater RDE should be, as shown in Fig. 4(e). And RDEs of real raining images could be found in Section V-B. Note that all RDEs are normalized between zero and one.

Note that the weight is updated adaptively according to Eq. (9). It implies that the rain density is estimated without any labels in our model, instead of training an extra network in a supervised way to classify the rain density into three categories in DID-MDN [8]. Besides, thanks to the weights, our model is capable of estimating the rain density with continuous states, which is more suitable for real raining scenes than algorithms only classifying rain density into limited discrete states, for instance, DID-MDN [8].

In this section, we will introduce the network structure of CODE-Net, mainly consisting of rain streaks extractor and rain streaks denoiser.

Rain Streaks Extractor. We use two convolutional layers $\boldsymbol{E}_{1}$, $\boldsymbol{E}_{2}$ to extract noisy rain streaks $\boldsymbol{r}_{\epsilon}$ from the rainy input $\boldsymbol{y}$:

as shown in Fig. 2, $\boldsymbol{r}_{\epsilon}$ can be seen as consisting of rain streaks and noise.

Rain Streaks Denoiser. $\boldsymbol{r}_{\epsilon}$ is then fed into Eq. (8) to seek for the convolutional sparse code $\boldsymbol{z}$. To avoid long inference time, we build Eq. (8) by CNN techniques.

Specifically, for $\Gamma_{\theta}(\cdot)$, under the assumption of nonnegative sparse coding, the soft nonnegative thresholding operator is equivalent to ReLU [21]. Thus, in Eq. (8), we replace $\Gamma_{\theta}\left(\alpha\right)$ with $\text{ReLU}\left(\alpha-\theta\right)$, where the threshold $\theta$ is learnable as well. Even though Eq. (12) is for one channel, the extension to multichannel case is mathematically straightforward. That is, for an input with $p$ channels, Eq. (12) still holds true with $\boldsymbol{G}\in\mathbb{R}^{c\times p\times s\times s}$, $\boldsymbol{S}\in\mathbb{R}^{c\times c\times s\times s}$, which means $\boldsymbol{S}$ and $\boldsymbol{G}$ correspond to two convolution layers. Above all, the resulting learning weighted ISTA (LwISTA) can be expressed as:

Based on Eq. (12), initializing $\boldsymbol{z}^{(0)}$ to $\boldsymbol{0}$, the optimal convolutional sparse code could be produced efficiently after $T$ iterations. Note that $\boldsymbol{S}$ shares the parameters over layers, as well as $\text{LWB}_{\boldsymbol{w}/\tilde{\boldsymbol{w}}}$.

After obtaining the convolutional sparse code $\boldsymbol{z}$, to recover the noiseless rain streak $\boldsymbol{r}$, another dictionary $\boldsymbol{E}_{3}$ corresponding to $\boldsymbol{G}$ and a convolutional layer $\boldsymbol{E}_{4}$ are needed:

Finally, the clean background image $\boldsymbol{x}$ would be produced by subtracting rain streaks $\boldsymbol{r}$ from the rainy image $\boldsymbol{y}$:

The whole architecture is illustrated in Fig. 2.

Datasets and Metrics. We use 12000 and 1800 pairs of images from [8, 6] and a high quality real rain dataset [7] as the training set. For testing, three commonly synthetic datasets, Rain1200 [8], Rain12 [2] and Test1000 [7], and some real-world images are utilized. The Rain1200 testing set, consisting of heavy, medium and light rain (Fig. 4(d)), is more challenging. Deraining results on synthetic datasets are evaluated with Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity (SSIM) [23] on the Y channel of transformed YCbCr space, while the performance of real-world images is evaluated visually since the ground truth images are not available.

Model. In CODE-Net, each convolutional layer has $128$ filters ($p=128$) of size $3\times 3$ ($s=3$) except $\boldsymbol{G}$, $\boldsymbol{S}$ and $\boldsymbol{E}_{4}$ respectively having $256$, $256$ ($c=256$) and $3$ filters. By contrast, we set $p=48$, $c=96$, and $s_{1}=3$, $s_{2}=5$, $s_{3}=7$ (namely, the rain filters are of size $3\times 3$, $5\times 5$ and $7\times 7$) in mCODE-Net¹¹1Note that, even though mCODE-Net is multiscale, it has fewer parameters ($802k$ vs. $1350k$) than CODE-Net since fewer filters are exploited.. We desert the biases and use zero-padding to keep the size of feature maps fixed. Another important point is, differ from [5, 8], there is no BatchNormalization [24] in our models according to Eq. (12) and Eq. (16).

Training. To make full use of the limited images, we augment the training set with random horizontal flips and $90$ degree rotations. The training process is divided into two stages: (1) we train two plain models (CODE-Net/mCODE-Net without weighting), initialized by the method of He et al. [25], using $\ell_{1}$ loss in RGB channels. For optimization, we use Adam [26] with $\beta_{1}=0.9$, $\beta_{2}=0.999$. In each training batch, 8 patches of size $128\times 128$ are extracted as inputs. The learning rate is set to $8\times 10^{-4}$ and halved at $[50k$, $100k$, $150k$, $200k]$ iterations; (2) we then initialize CODE-Net and mCODE-Net using the plain models and fine-tune them with the learning rate set to $8\times 10^{-5}$ and halved at $[50k$, $100k$, $150k$, $200k]$ iterations. We use PyTorch to implement our models with an NVIDIA RTX 2080 GPU²²2Codes and more results are available at https://github.com/Achhhe/CODE-Net.

In this subsection, we focus on the rain density estimation of rainy images with synthetic and real rain. For synthetic rainy images, we show a sample of different rain levels in Fig. 4(d). And in Fig. 4(e), RDEs of several samples are depicted. One can see that, the RDEs mainly range from $0.1$ to $0.95$ and become larger as the rain density increases, consistent with the analysis in Section III-B.

Besides, a clear image and some real-world rainy images with continuous rainy states and their RDEs are shown in Fig. 5. The clear image has the smallest RDE $0.1617$ while the rainy images have gradual increased RDEs from $0.2550$ to $0.9850$, which verifies the validity of continuous rain density estimation of our model.

In this subsection, we compare our proposed CODE-Net and mCODE-Net to 9 state-of-the-art deraining methods, including DSC [1], GMM [2], CNN [4], DDN [5], JORDER [6], DID-MDN [8], SEMI [18], SPANet [7] and PreNet [19]. All results are reproduced by the authors’ codes except the quantitative ones of DSC and GMM on Rain1200 are cited from [8] directly since their codes run slowly.

Synthetic Images. Quantitative results are tabulated in Tab. I. Clearly, both CODE-Net and mCODE-Net achieve superior performance. Especially on Rain1200, the more complex dataset consisting of light, medium and heavy levels, our models outperform others by large margins, which indicates our models are more robust for different densities of rain. One important reason is that the rain density is implicitly considered in Eq. (8) and thus leads to better rain removal results. Fig. 6 and Fig. 7 show the visual comparisons. Our models consistently achieve the best visual performance in terms of effectively removing rain streaks while preserving background details. Besides, thanks to the multiscale prior, mCODE-Net could achieve comparable or even better results than CODE-Net with fewer parameters.

Note that, JORDER and DID-MDN use additional rain mask or density information in training stage. Nonetheless, our models still greatly outperform their results without any additional data.

Real Images. To demonstrate the generalization of the proposed architectures, we evaluate our models and other methods on four real-world images, shown in Fig. 8 and Fig. 9. Once again, our models show the best visual performance. In contrast, previous methods especially DID-MDN tend to produce under/over deraining results. Our models get rid of these faults and produce more pleasant results, which verifies the effectiveness of continuous rain density estimation over classifying it into three levels.

Comparing CODE-Net and mCODE-Net, raindrops with large or small size can be efficiently removed by mCODE-Net with further consideration of the multi-scale property, while CODE-Net gives a result by leveraging different raindrop scales, as shown in Fig. 8.

We further apply the CODE-Net on the deraining results obtained by each method and use RDE Eq. (10) to estimate their rain densities. As shown in Fig. 8 and Fig. 9, the RDEs of the deraining images decrease comparing to the original rainy inputs, which reflects the deraining effects of each method. One can find that the proposed CODE-Net and mCODE-Net can achieve much lower RDEs than other methods for rainy images with heavy rains shown in Fig. 8. It coincides with the results on the synthetic rainy images as shown in Tab. I, where CODE-Net and mCODE-Net can outperform the state-of-the-arts by large margins.

On the other hand, the RDEs obtained from CODE-Net have potentials to quantitatively evaluate the deraining performance for rainy images from real scenes.

We visualize the running time per image ($512\times 512$) versus performance (PSNR) on Rain1200 dataset in Fig. 10. According to the provided source code, all methods are tested on GPU except DSC and GMM on CPU. One key factor affecting running time of our method is the number of iterations $T$ when solving convolutional sparse codes in Eq. (12). For simplicity, we denote CODE-Net/mCODE-Net iterating $T$ times as C_T/mC_T. It is observed that C₅ achieves comparable performance and computational complexity as other state-of-the-art methods. However, too few iterations are not capable to generate good convolutional sparse codes. Hence, we additionally test three models with $T=15,25,35$ and unsurprisingly the results of them outperform those of other methods largely.

As shown in Fig. 10, with the increase of the number of iterations, the performance of our method can be further improved and obtain the optimal result when $T=25$. In this paper, we focus on the best rain removal performance and thus choose $T=25$ in our final models (C₂₅/mC₂₅). For practical applications, we could adjust the number of iterations $T$ to balance performance and running time.

Our network (Fig. 2) is derived from traditional CSC problem, but by coincidence, it includes two common strategies in deep learning, i.e., local residual learning (LRL) [27] and pre-activation (PA) [28]. LRL is actually the local skip connection and PA means activation layer (LWB-ReLU-LWB) comes before weight layer ($\boldsymbol{S}$) in unfolded LwISTA. Besides, to improve the deraining performance, we introduce reweighting (RW) into CSC optimization. Last but not the least, since their specific characteristics, we focus on extracting rain streaks and recover the deraining image by subtracting them from the input. Interestingly, this corresponds to the global residual learning (GRL).

To demonstrate the effectiveness of aforementioned components, we conduct an ablation study of GRL, LRL, PA and RW on CODE-Net. $96$ images are selected from Rain1200 for testing, with each level $32$ images. Tab. II shows the average PSNR. When both LRL and PA are used, the PSNR is relatively high, no matter RW is used or not (\nth2/\nth3 vs. \nth8/\nth9). Comparing \nth2/\nth4/\nth6/\nth8 with \nth3/\nth5/\nth7/\nth9, one can see that the performance could be improved by adding RW. However, it is interesting that, using PA only would degrade the performance (\nth2/\nth3 vs. \nth4/\nth5). Not surprisingly, CODE-Net using four components simultaneously performs the best. These comparisons indicate the superiority of reweighted CSC inspired network combining GRL, LRL, PA and RW in a particular way.

Though the proposed network could achieve superior performance on most of testing images while producing corresponding RDEs, some situations occurred in which our method would produce overestimation of RDE and over deraining result, as exhibited in Fig. 11. This is primarily caused by the confusing of rain-like structures and the crude global average pooling operation in LWB Eq. (9) that only considers the overall information of each channle. We believe that the predicament could be tackled if we capture sufficient local details additionally.

In this section, we explore the potentials of our network for other visual tasks. In the pursuit of better deraining effect, we first train a generative adversarial CODE-Net (GaCODE-Net) using perceptual loss [29, 30]. As shown in Fig. 13(e), although inferior to CODE-Net on quantitative metrics (PSNR and SSIM), GaCODE-Net is capable of generating more realistic textures and high-frequency details. Moreover, since snowy image could be modeled as Eq. (1) [31] and snow streaks share much similarity with rain streaks, we then evaluate our CODE-Net on image desnowing task. Fig. 13(h) shows the results on two real snowy images, which verifies the promising generalization of our method on other tasks with similar degeneration model.