Hong Wang, Qi Xie, Qian Zhao, and Deyu Meng
[PDF] [Supplementary Material]
Deep learning (DL) methods have achieved state-of-the-art performance in the task of single image rain removal. Most of current DL architectures, however, are still lack of sufficient interpretability and not fully integrated with physical structures inside general rain streaks. To this issue, in this paper, we propose a model-driven deep neural network for the task, with fully interpretable network structures. Specifically, based on the convolutional dictionary learning mechanism for representing rain, we propose a novel single image deraining model and utilize the proximal gradient descent technique to design an iterative algorithm only containing simple operators for solving the model. Such a simple implementation scheme facilitates us to unfold it into a new deep network architecture, called rain convolutional dictionary network (RCDNet), with almost every network module one-to-one corresponding to each operation involved in the algorithm. By end-to-end training the proposed RCDNet, all the rain kernels and proximal operators can be automatically extracted, faithfully characterizing the features of both rain and clean background layers, and thus naturally lead to its better deraining performance, especially in real scenarios. Comprehensive experiments substantiate the superiority of the proposed network, especially its well generality to diverse testing scenarios and good interpretability for all its modules, as compared with state-of-the-arts both visually and quantitatively.
@InProceedings{Wang_2020_CVPR,
author = {Wang, Hong and Xie, Qi and Zhao, Qian and Meng, Deyu},
title = {A Model-Driven Deep Neural Network for Single Image Rain Removal},
booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2020}
}
- Linux Platform
- NVIDIA GPU + CUDA CuDNN (CUDA 10.0)
- PyTorch == 0.4.1 (1.0.x and higher version need revise the dataloader.py in folder "../src/", refer to "Friendly Tips" for details)
- Python3.6
- imageio2.5
- numpy1.14
- opencv-python
- scikit-image0.13.0
- tqdm4.32.2
- Rain100L: 200 training pairs and 100 testing pairs
- Rain100H: 1800 training pairs and 100 testing pairs
- Rain1400: 12600 training pairs and 1400 testing pairs
- Rain12: 12 testing pairs
- 638492 training pairs and 1000 testing pairs
- 147 unlabeled rainy images from SIRR
*More detailed explanations refer to the single image part in the Survey
taking training Rain100L as an example:
- Download Rain100L (including training set and testing set) from the [NetDisk] (pwd:uz8h) and put them into the folder "./RCDNet_code/for_syn/data", then the content is just like:
"./for_syn/data/train/small/rain/*.png"
"./for_syn/data/train/small/norain/*.png"
"./for_syn/data/test/small/rain/*.png"
"./for_syn/data/test/small/norain/*.png"
- Begining training:
$ cd ./RCDNet_code/for_syn/src/
$ python main.py --save RCDNet_syn --model RCDNet --scale 2 --epochs 100 --batch_size 16 --patch_size 64 --data_train RainHeavy --n_threads 0 --data_test RainHeavyTest --stage 17 --lr_decay 25 --gamma 0.2 --num_M 32 --num_Z 32 --data_range 1-200/1-100 --loss 1*MSE --save_models
*Note that: --data_range 1-#training pairs/1-#testing pairs. The command above is set based on Rain100L (200 training paris and 100 testing pairs).
- Download the training set from [SPANet] and testing set from the [NetDisk] (pwd:uz8h) and put them into the folder "./RCDNet_code/for_spa/data", then the content is just like:
"./for_spa/data/real_world"
"./for_spa/data/real_world_gt"
"./for_spa/data/real_world.txt"
"./for_spa/data/test/small/rain/*.png"
"./for_spa/data/test/small/norain/*.png"
*Note that: download the training set from [SPANet] , which includes three files: "real_world"， "real_world_gt", and "real_world.txt"
- Begining training:
$ cd ./RCDNet_code/for_spa/src/
$ python main.py --save RCDNet_spa --model RCDNet --scale 2 --epochs 100 --batch_size 16 --patch_size 64 --data_train RainHeavy --n_threads 0 --data_test RainHeavyTest --stage 17 --lr_decay 25 --gamma 0.2 --num_M 32 --num_Z 32 --data_range 1-638492/1-1000 --loss 1*MSE --save_models
$ cd ./RCDNet_code/for_syn/src/
$ python main.py --data_test RainHeavyTest --ext img --scale 2 --data_range 1-200/1-100 --pre_train ../experiment/RCDNet_syn/model/model_best.pt --model RCDNet --test_only --save_results --save RCDNet_test
*Note that: --data_range 1-#training pairs/1-#testing pairs. The command (-- data_range 1-200/1-100) above is set based on Rain100L (200 training paris and 100 testing pairs).
The derained results are saved in the folder "./for_syn/experiment/RCDNet_test/results/", where the image name "norain-*LR.png", "norain-*HR.png" , "norain-*SR.png" means rainy image, groundtruth, and restored background, respectively.
$ cd ./RCDNet_code/for_spa/src/
$ python main.py --data_test RainHeavyTest --ext img --scale 2 --data_range 1-638492/1-1000 --pre_train ../experiment/RCDNet_spa/model/model_best.pt --model RCDNet --test_only --save_results --save RCDNet_test
The derained results are saved in the folder "./for_spa/experiment/RCDNet_test/results/"
See folder "Pretrained_model"
-
for synthetic model: put the pretrained model into the folder "./for_syn/experiment/RCDNet_syn/model/", and then the content is just like: "../for_syn/experiment/RCDNet_syn/model/model_best.pt"
-
for spa model: put the pretrained model into the folder "./for_spa/experiment/RCDNet_spa/model/", and then the content is just like: "../for_spa/experiment/RCDNet_spa/model/model_best.pt"
If needed, please download from [NetDisk](pwd:e2xo)
All PSNR and SSIM results are computed by using this Matlab code, based on Y channel of YCbCr space.
Code borrows from JORDER_E by Wenhan Yang. Thanks for sharing !
- If higher version is needed, please replace the files "dataloader.py" and "trainer.py" in the folder "./RCDNet_code/../src/" with the corresponding ones in the folder "./pytorch1.0+/../src/" .
- More general and friendly code framework for RCDNet will be coming !
2020.6.10 Release code
2020.6.22 Fix the bug in the file "../src/model/rcdnet.py", change the training command, and also release higher version with the folder "./pytorch1.0+"
2020.6.28 Fix the bug in the file "init_rain_kernel.m"
If you have any question, please feel free to concat Hong Wang (Email: hongwang01@stu.xjtu.edu.cn)