MCGKT-Net: Multi-level Context Gating Knowledge Transfer Network for Single Image Deraining

Rain streaks removal from a single image is an extremely challenging task due to its ill-posed nature. Previous methods are mainly divided into two categories: optimization-based methods [17, 15, 18, 25] and deep learning-based methods [19, 20, 31, 21]. Optimization based methods usually formulate the deraining task as a mathematical model motivated by the fact that rainy images can be decomposed into a clean background image layer and a rain layer. To recover more robust clean image, the prior knowledge for characterizing the underlying structure of the latent clean image layer and the attribute of the rain layers has imposed on the formulated mathematical model as regularization term and employed optimization strategy for solving. Kang et al. [15] proposed to apply sparse coding to separate rain streaks from the high-frequency layer, while Luo et al. [16] explored a discriminative sparse coding framework for modeling image patches. The work by Chen et al. [17] and Chang et al. [18] leveraged the low-rank property of rain streaks for removing the decomposed rain layer based on low-rankness. Since the composite models [17, 15, 18] regularized by modeling the sparse and low-rank prior are insufficient in characterizing the decomposed layers, leading to limited deraining performance on diverse images. In addition, the explored priors (for example sparsity, low-rankness) on the previous approaches are hand-crafted, and to discover a proper prior for a specific image and a kind of rain type still remains be an art or requires comprehensive analysis for a specific rainy image. Recently, deep convolutional neural network has been widely applied to single image deraining, and validated that promising performance can be achieved [19, 20, 31, 21]. Fu et al. [19] first explored a three-layer convolutional network to predict clean image high-frequency component from its rain-contaminated counterpart, and further extended it to a 20-layer CNN structure by incorporating Residual-Block, called as deep detail network, for pursuing better performance [20]. Zhang et al. [31] presented a multi-stream dense network for joint rain density estimation and deraining. To generate more visually plausible deraining result, the same research group investigated a conditional generative adversarial network (GAN) for single image deraining [21], and proved to achieve visually high-quality reconstructions. In [36], a novel deep network architecture based on recurrent neural networks and squeeze-and-excitation context aggregation module (RESCAN) has been proposed and adaptively adjusted parameters for various rain streak layers. Fan et al. [23] proposed residual-guide network with recursive convolution module and multi-level supervision not only on the final results but also on the intermediate results progressively for predicting high-quality reconstruction. Wei et al. [30] proposed a semi-supervised image deraining network, while Ren et al. [37] focused on several factors including network architecture for integrating progressive ResNet and recurrent layers inside and cross stages, and loss functions, and provided a better and simpler baseline deraining network.

It is known that a rainy image is possibly decomposed into the image layer and rain streak layer, which may consist of multiple layers, especially under heavy rain conditions. The rain streaks decomposed in multiple layers manifest intrinsically multi-scale attributes and some self-similarity properties within and across scales, which is prospected to boost the deraining performance via exploring the correlated information between and across multiple levels of rain layers. Most existing methods recur to deeper and more complex network architecture for pursuing better deraining performance but cannot make full use of the underlying correlation between and across different rain layers. Although, a few work [34, 8, 9] have been investigated to explore multiscale information for deraining from a single image, which mainly leverages multiple subnets (several mainstreams) for exploiting different scales, and lead to more complicated DCNN architectures. As we know that the convolutional encoder-decoder network itself is a multi-level learning architecture, where the encoder path learns feature representation evolved to large scale with the increased depth of the network while the decoder path attempts to recover the feature representation with more detailed structure (small scale) from the final output of the encoder with more semantic information and large context. Further, to retain more detail structures in the final prediction results, skip connection is usually used for duplicating the feature representation of the encoder path to the decoder side in the same level such as in U-Net [10], FCN [11]. Although the encoder-decoder network has also been adopted for image deraining [12, 13, 14], most existing methods cannot effectively exploit the correlation of the feature representations and compensate each other between the encoder and decoder paths. This study explores a novel and simple deraining network, called multi-level context gating knowledge transfer network (MCGKT-Net), which is based on the multi-level encoder-decoder network and investigates both internal and external knowledge transfer for boosting deraining performance.

The mainstream of our used network architecture consists of two paths: encoder and decoder, and each path is divided into four blocks. Both encoder and decoder paths learn multi-level feature representations in the multiple blocks, where encoder employs MaxPooling layer with a 2*2 kernel for decreasing feature map size to half in both horizontal and vertical directions between blocks while decoder performs up-sampling for doubly recovering the feature map size between blocks. In each block of both encode and decoder, we implement it in 3 convolutional layers with 3*3 kernels following ReLU activation function. The channel number of the learned feature maps block-wisely is doubled in the encoder while is halve in the decoder. Thus the architecture in the encoder path is same as the first 3 shallow layers of the popularly used VGGNet in different vision problems.

Let’s denotes the input and output of the $j-th$ block in the encoder path as ${\bf E}_{i}^{j}$, ${\bf E}_{o}^{j}$, and in the decoder path as ${\bf D}_{i}^{j}$, ${\bf D}_{o}^{j}$, respectively, the relation of the input and output of each block can be expressed as:

$${\bf E}_{o}^{j}=f({\bf E}_{o}^{j},\theta_{E,i}),\;{\bf D}_{o}^{j}=f({\bf D}_{o}^{j},\theta_{D,i})$$

(2)

where $f(\cdot)$ represents the transformation operation of 3 convolutional layers with the learned parameters $\theta_{E,i}$ and $\theta_{D,i}$, respectively. The $j+1-th$ inputs in the encoder and the $j-1-th$ output in the decoder sides can be obtained from the $j-th$ outputs, and are expressed as:

$${\bf E}_{i}^{j+1}=MP({\bf E}_{o}^{j}),\;{\bf D}_{i}^{j-1}=UP({\bf D}_{o}^{j})$$

(3)

where $MP(\cdot)$ and $UP(\cdot)$ denote the MaxPooling and up-sampling layer implemented between blocks, respectively. Then the mainstream of our encoder-decoder network can learn multi-level feature representations: ${\bf E}_{i}=[{\bf E}_{i}^{1},{\bf E}_{i}^{2},{\bf E}_{i}^{3},{\bf E}_{i}^{4}]$ in the encoder and ${\bf D}_{o}=[{\bf D}_{o}^{1},{\bf D}_{o}^{2},{\bf D}_{o}^{3},{\bf D}_{o}^{4}]$ in the decoder with different context information, where the encoder generates low-level features with more detail image structure while decoder results in high-level feature representation with more semantic information. The encoder-decoder architecture is very simple with the convolutional layer of small kernel size 3*3 and is prospected to be easily trained in different vision tasks. In this network, it is known that the outputs of the deeper blocks inherit from the shallow blocks of the decoder and all blocks of the encoder, and produce more semantic information but may lose detail structure. To maintain more detail structure in the final results, the existing encoder-decoder architecture in different vision tasks such as FCN, U-Net usually employ skip connection for directly duplicating the output of the encoder to the corresponding input sides in the decoder, which cannot effectively leverage the learned feature maps in both encoder and decoder paths. This study integrates several elementary modules for effectively interaction learning or knowledge transfer among the feature representations of the main network and adaptively selects more useful information. Next, we describe the integrated modules: knowledge transfer module and multi-level context gating module.

As described above, the simple feature reuse via skip connection cannot effectively fuse the learned feature representations in both encoder and decoder paths. In addition, it is known that the feature representations with more semantic information are gradually learned based on the former ones and inherit from the already learned features including those in the encoder. However, the feature representations with more semantic information in the decoder path cannot be back transferred to the encoder side for calibrating the low-level features. This study integrates a transfer module among the corresponding blocks in the encoder and decoder paths and conducts back transfer, called as internal knowledge transfer (IKT) module. Furthermore, we also aim at investigating the already learned knowledge in the released CNN models in other vision domains for aiding our deraining network training, called as external knowledge transfer (EKT) module.

It is obvious that the network can obtain large amount of feature representations in different layers of multi-level blocks. However, not all learned features equivalently contribute to the subsequent extraction of high representative features and the final prediction. Recently, attention mechanism has been popularly explored to adaptively concentrate the more discriminated and effective features in the network training procedure. For example, motivated by the fact that different regions in the input may have various contributions to the final prediction such as in image classification, detection and segmentation, many work exploit the spatial attention via mutual enhancement with spatial correlation, and manifest significant improvement. Our goal aims at predicting all underlying pixel values from the rain-degraded input, and all regions should be indispensable for estimating the precise pixel values nearby. Thus, this study instead investigates the channel attention for dedicating to emphasize the channel with the underlying scene information. We suppose that the feature maps extracted by various convolutional kernels may correspond to some underlying scene layers or a part of rain layers, and propose to exploit explicit relationship between channels of the convolutional layers for gating context. We implement the context gating module via adaptively assigning a weight for each channel (channel attention) and then encoding the inputted raw feature maps.

The multi-level context gating (MLCG) module shown in Fig. 2, consists of two-part: squeeze and excitation (also called SE-block), and is employed on multi-level integrated feature maps [${\bf\hat{D}}_{i}^{1}$, ${\bf\hat{D}}_{i}^{2}$ , $\cdots$, ${\bf\hat{D}}_{i}^{L}$] of the encoder and decoder paths, which are also the input of the decoder’s blocks. The input feature maps to MLCG module are aggregated to generate channel contribution index by employing global average pooling (GAP) of the whole context of channels. Let’s denote the input feature map to the $l-th$ MLCG module as ${\bf\hat{D}}_{i}^{l}=[{\bf\hat{d}}_{i,1}^{l},{\bf\hat{d}}_{i,2}^{l}$, $\cdots$, ${\bf\hat{d}}_{i,C}^{l}]$, where ${\bf\hat{d}}_{i,c}^{l}\in\mathbb{R}^{W\times H}$, and the spatial squeeze (GAP) is formulated as:

$$s_{c}=f_{sq}({\bf\hat{d}}_{i,c}^{l})=\frac{1}{W\times H}\sum_{h}^{H}\sum_{w}^{W}{\bf\hat{d}}_{i,c}^{l}(h,w)$$

(5)

where $f_{sq}(\cdot)$ is the spatial squeeze function for compressing each two-dimensional feature map as a contribution index $s_{c}$, and $W\times H$ is the size of the $c-th$ channel feature map. Then, we employ excitation functions for capturing the channel-wise dependencies and non-linear interaction based on the global channel information ${\bf s}=[s_{1},s_{2},\cdots,s_{C}]$, which is implemented with two fully connected (FC) layers in the MLCG module. The first FC layer encodes the channel global vector ${\bf s}$ to a dimension-reduced vector with reduction ratio $r$, and the second FC layer encodes it back again to the dimension $C$ as an excitation vector, which can be expressed:

$$z=f_{ex}({\bf s},{\bf W})=\delta({\bf W}_{2}\sigma({\bf W}_{1}{\bf s}))$$

(6)

where ${\bf W}_{1}\in\mathbb{R}^{\frac{C}{r}\times C}$ and ${\bf W}_{2}\in\mathbb{R}^{C\times\frac{C}{r}}$ are the parameters of the first and second FC layers, respectively, $\sigma(\cdot)$ and $\delta(\cdot)$ refer ReLU and sigmoid activation functions. The final output of the MLCG module is generated as:

$${\bf\tilde{d}}_{i,c}^{l}=f_{scale}({\bf\hat{d}}_{i,c}^{l},z_{c})=z_{c}{\bf\hat{d}}_{i,c}^{l}$$

(7)

where $f_{scale}$ denotes a channel-wise multiplication between the channel attention index $z_{c}$ and the input feature map.

We compare our proposed MCGKT-Net with the state-of-the-art deraining approaches, including semi-supervised transfer learning (SEMI) [30], density-aware deraining (DIDMDN) [31], simple deep convolutional network for image SR (SRCNN) [32], deep detail network (DDN) [20], image-to-image translation (pix2pix)  [33], lightweight pyramid network (LPNet) [34], U-Net [10], uncertainty guided multi-scale residual learning (UMRL) [35], recurrent squeeze-and-excitation context aggregation net (RESCAN) [36], progressive deraining network (PreNet) [37]. Table 1 shows the quantitative measure on the three datasets. From Table 1, we can see that our methods can outperform almost methods. For providing visual comparison, Fig. 4 and 5 visualize the derained examples from Rain100L and Rain100H datasets using different methods, which manifests that our proposed MCGKT-Net can recover more clean images than other existing methods. Furthermore, we also provide the derained results on three real images in Fig. 6 using our proposed MCGKT-Net and several state-of-the-art methods.

Our proposed MCGKT-Net is evolved from the baseline U-Net architecture, which is very simple and easy to be trained effectively. We integrated three modules: internal knowledge transfer (IKT), external knowledge transfer (EKT), and multi-level context gating (MLCG) modules into the baseline U-Net for higher representative feature learning. Therein EKT module adopted the pre-trained VGG models: VGG16 and VGG19. This section evaluates the effectiveness of the integrated modules on the baseline U-Net. Quantitative results on all three datasets are given in Table 2. It can be seen from Table 2 that the proposed MCGKT-Net manifests great superiority over the baseline U-Net and its incomplete versions not integrating all proposed modules. The results of the experiment with the addition of each of the three modules in the Rain800 are shown in Fig. 7. Integration of all modules surpasses the baseline by 3.78dB, 4.26dB, 1.16dB for Rain100H, Rain100L, and Rain800, respectively.

From the results of Table 2, all integrated modules can improve the quantitative measures on both Rain100H and Rain100L datasets while the integration of the MLCG module on the Rain800 dataset decreases the quantitative metrics a little. From Fig. 3 as mentioned above that the rain streaks in Rain100H and Rain100L datasets exhibit some regular characteristics similar to the line structures with diverse directions while the rain streaks in the Rain800 dataset have no regular pattern most like un-regular noise. The intent of integrating the MLCG module is to adaptively emphasize and attenuate specific channels of features with some specific patterns and is expected to be oriented well to the existed rain steaks in the Rain100H and Rain100L datasets while be difficult to attenuate the noise-like rain streaks in the Rain800 dataset.