From Rain Generation to Rain Removal (CVPR2021)

Hong Wang, Zongsheng Yue, Qi Xie, Qian Zhao, Yefeng Zheng, and Deyu Meng

Abstract

For the single image rain removal (SIRR) task, the performance of deep learning (DL)-based methods is mainly affected by the designed deraining models and training datasets. Most of current state-of-the-art focus on constructing powerful deep models to obtain better deraining results. In this paper, to further improve the deraining performance, we novelly attempt to handle the SIRR task from the perspective of training datasets by exploring a more efficient way to synthesize rainy images. Specifically, we build a full Bayesian generative model for rainy image where the rain layer is parameterized as a generator with the input as some latent variables representing the physical structural rain factors, e.g., direction, scale, and thickness. To solve this model, we employ the variational inference framework to approximate the expected statistical distribution of rainy image in a data-driven manner. With the learned generator, we can automatically and sufficiently generate diverse and non-repetitive training pairs so as to efficiently enrich and augment the existing benchmark datasets. User study qualitatively and quantitatively evaluates the realism of generated rainy images. Comprehensive experiments substantiate that the proposed model can faithfully extract the complex rain distribution that not only helps significantly improve the deraining performance of current deep single image derainers, but also largely loosens the requirement of large training sample pre-collection for the SIRR task.

Dependicies

pip install -r requirements.txt

Folder Directory

.
|-- for_spa                                   : Experiments on real SPA-Data
|   |-- data                                  : SPA-Data: train + test
|   |   `-- spa-data 
|   |       |-- real_world              
|   |       |-- real_world.txt
|   |       |-- real_world_gt
|   |       `-- test  
|   |-- train_spa_joint.py                    : Joint training on SPA-Data
|   |-- train_spa_aug.py                      : Augmentated training
|   |-- train_spa_smallsample_aug.py          : Small sample experiments (GNet in Table 1)
|   |-- train_spa_smallsample_noaug.py        : Small sample experiments (Baseline in Table 1)
|   |-- test_disentanglement.py               : Distentanglement experiments on SPA-Data
|   |-- test_interpolation.py                 : Interpolation experiments on SPA-Data
|   |-- spamodels                             : Joint pretrained model on SPA-Data

|-- for_syn                                   : Experiments on synthesized datasets
|   |-- data                                  : Synthesized datasets: train + test
|   |   |-- rain100H
|   |   |   |-- test
|   |   |   `-- train
|   |   |-- rain100L
|   |   |   |-- test
|   |   |   `-- train
|   |   `-- rain1400
|   |       |-- test
|   |       `-- train
|   |-- train_syn_joint.py                    : Joint training
|   |-- train_syn_aug.py                      : Augmentated training in Table 2
|   |-- test_disentanglement.py               : Distentanglement experiments
|   |-- test_interpolation.py                 : Interpolation experiments 
|   |-- syn100hmodels                         : Joint pretrained model on rain100H
|   |-- syn100lmodels                         : Joint pretrained model on rain100L
|   |-- syn1400models                         : Joint pretrained model on rain1400

Benchmark Dataset

Synthetic datasets: Rain100L, Rain100H, Rain1400

Real datasets: SPA-Data, Internet-Data（only for testing）

Detailed descriptions refer to the Survey, SCIENCE CHINA Information Sciences2021

Please refer to RCDNet, CVPR2021 for downloading these datasets and put them into the corresponding folders according to the dictionary above.

For Synthetic Dataset (taking Rain100L as an example)

Training

Step 1. Joint Training:

$ cd ./VRGNet/for_syn/ 
$ python train_syn_joint.py  --data_path "./data/rain100L/train/small/rain" --gt_path "./data/rain100L/train/small/norain" --log_dir "./syn100llogs/" --model_dir "./syn100lmodels/" --gpu_id 0

Step 2. Augmentated Training: (taking baseline PReNet as an example)

$ python train_syn_aug.py  --data_path "./data/rain100L/train/small/rain" --gt_path "./data/rain100L/train/small/norain" --netED "./syn100lmodels/ED_state_700.pt" --log_dir "./aug_syn100llogs/" --model_dir "./aug_syn100lmodels/" --fake_ratio 0.5 --niter 200 --gpu_id 0

Testing

Joint Testing:

$ python test_syn_joint.py  --data_path "./data/rain100L/test/small/rain" --netDerain "./syn100lmodels/DerainNet_state_700.pt" --save_path "./derained_results/rain100L/" --gpu_id 0

Augmentated Testing: (taking baseline PReNet as an example)

$ python test_syn_aug.py  --data_path "./data/rain100L/test/small/rain" --model_dir "./aug_syn100lmodels/Aug_DerainNet_state_200.pt" --save_path "./aug_derained_results/rain100L/" --gpu_id 0

Interpolation Testing:

$ python test_interpolation.py   --data_path "./interpolation_results/test_data/rain100L/rain" --gt_path "./interpolation_results/test_data/rain100L/norain" --netED "./syn100lmodels/ED_state_700.pt"  --save_patch "./interpolation_results/test_data/rain100L/crop_patch/" --save_inputfake "./interpolation_results/generated_data/rain100L/input_fake" --save_rainfake "./interpolation_results/generated_data/rain100L/rain_fake" --gpu_id 0

Disentanglement Testing:

$ python test_disentanglement.py  --netED "./syn100lmodels/ED_state_700.pt" --save_fake "./disentanglement_results/rain100L/" --gpu_id 0

For SPA-Data

Training

Step 1. Joint Training:

$ cd ./VRGNet/for_spa/ 
$ python train_spa_joint.py  --data_path "./data/spa-data/" --log_dir "./spalogs/" --model_dir "./spamodels/" --gpu_id 0

Step 2. Augmentated Training: (taking baseline PReNet as an example)

$ python train_spa_aug.py  --data_path "./data/spa-data/" --netED "./spamodels/ED_state_800.pt" --log_dir "./aug_spalogs/" --model_dir "./aug_spamodels/" --fake_ratio 0.5 --niter 200 --gpu_id 0

Step 3. Small Sample Training: (taking baseline PReNet as an example)

$ python train_spa_smallsample_aug.py  --data_path "./data/spa-data/" --netED "./spamodels/ED_state_800.pt" --fake_ratio 0.5 --train_num 1000 --log_dir "./aug05_spalogs/" --model_dir "./aug05_spamodels/" --niter 200 --gpu_id 0

$ python train_spa_smallsample_noaug.py  --data_path "./data/spa-data/" --fake_ratio 0.5 --train_num 1000 --log_dir "./noaug05_spalogs/" --model_dir "./noaug05_spamodels/" --niter 200 --gpu_id 0

Testing

Joint Testing:

$ python test_spa_joint.py  --data_path "./data/spa-data/test/small/rain" --netDerain "./spamodels/DerainNet_state_800.pt" --save_path "./derained_results/spa-data/" --gpu_id 0

Augmentated Testing: (taking baseline PReNet as an example)

$ python test_spa_aug.py  --data_path "./data/spa-data/test/small/rain" --model_dir "./aug_spamodels/Aug_DerainNet_state_200.pt" --save_path "./aug_derained_results/spa-data/" --gpu_id 0

Interpolation Testing:

$ python test_interpolation.py   --data_path "./interpolation_results/test_data/spa-data/rain" --gt_path "./interpolation_results/test_data/spa-data/norain" --netED "./spamodels/ED_state_800.pt"  --save_patch "./interpolation_results/test_data/spa-data/crop_patch/" --save_inputfake "./interpolation_results/generated_data/spa-data/input_fake" --save_rainfake "./interpolation_results/generated_data/spa-data/rain_fake" --gpu_id 0

Disentanglement Testing:

$ python test_disentanglement.py  --netED "./spamodels/ED_state_800.pt" --save_fake "./disentanglement_results/spa-data/" --gpu_id 0

Small Sample Testing: (taking baseline PReNet as an example)

$ python test_spa_aug.py  --data_path "./data/spa-data/test/small/rain" --model_dir "./aug05_spamodels/Aug05_DerainNet_state_200.pt" --save_path "./aug05_derained_results/spa-data/" --gpu_id 0

$ python test_spa_aug.py  --data_path "./data/spa-data/test/small/rain" --model_dir "./noaug05_spamodels/NoAug05_DerainNet_state_200.pt" --save_path "./noaug05_derained_results/spa-data/" --gpu_id 0

For Internet-Data

The test model is trained on SPA-Data.

Pretrained Model and Usage

We have provided the joint pretrained model saved in syn100lmodels, syn100hmodels, syn1400models, and spamodels. If needed, you can dirctly utilize them to augment the original training set without exectuting the joint training.
We only provide the PReNet for an example during the augmented training/testing phase. This is a demo. In practice, you can easily replace PReNet with other deep deraining models as well as yours for further performance improvement by adopting the augmented strategy with our generator. Please note that the training details in train_syn_aug.pyand train_spa_aug.pyare needed to be correspondingly adjusted.
Please note that in our default settings, the generated patchsize is 64x64. In the released code, we also provide the model revision (i.e., RNet, Generator, and discriminator) for generating the size as 256x256. If other sizes are needed, you can correspondingly revise the network layer and then re-train the joint VRGNet.

Rain Generation Experiments

Rain Removal Experiments

Derained Results of Our VRGNet (i.e., PReNet-)

All PSNR and SSIM results are computed with this Matlab code. If needed, please download the results from NetDisk (pwd:2q6l)

Citation

@InProceedings{Wang_2021_CVPR,  
author = {Wang, Hong and Yue, Zongsheng and Xie, Qi and Zhao, Qian and Zheng, Yefeng and Meng, Deyu},  
title = {From Rain Generation to Rain Removal},  
booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},  
month = {June},  
year = {2021}  
}

Contact

If you have any question, please feel free to concat Hong Wang (Email: [email protected])

【CVPR 2021, Variational Inference Framework, PyTorch】 From Rain Generation to Rain Removal

Related tags

Overview

From Rain Generation to Rain Removal (CVPR2021)

Abstract

Dependicies

Folder Directory

Benchmark Dataset

For Synthetic Dataset (taking Rain100L as an example)

Training

Testing

For SPA-Data

Training

Testing

For Internet-Data

Pretrained Model and Usage

Rain Generation Experiments

Rain Removal Experiments

Derained Results of Our VRGNet (i.e., PReNet-)

Citation

Contact

Owner

Hong Wang

Person Re-identification

A simple image/video to Desmos graph converter run locally

[Nature Machine Intelligence' 21] "Advancing COVID-19 Diagnosis with Privacy-Preserving Collaboration in Artificial Intelligence"

Deep Learning as a Cloud API Service.

This repository contains the map content ontology used in narrative cartography

Deep Surface Reconstruction from Point Clouds with Visibility Information

You Only 👀 One Sequence

Code for the paper 'A High Performance CRF Model for Clothes Parsing'.

[NeurIPS 2021] Large Scale Learning on Non-Homophilous Graphs: New Benchmarks and Strong Simple Methods

Dynamic Attentive Graph Learning for Image Restoration, ICCV2021 [PyTorch Code]

MOT-Tracking-by-Detection-Pipeline - For Tracking-by-Detection format MOT (Multi Object Tracking), is it a framework that separates Detection and Tracking processes?

Pairwise learning neural link prediction for ogb link prediction

Convolutional Neural Networks

The Rich Get Richer: Disparate Impact of Semi-Supervised Learning

Activating More Pixels in Image Super-Resolution Transformer

Chainer Implementation of Fully Convolutional Networks. (Training code to reproduce the original result is available.)

Prompt-BERT: Prompt makes BERT Better at Sentence Embeddings

Scale-aware Automatic Augmentation for Object Detection (CVPR 2021)

Home for cuQuantum Python & NVIDIA cuQuantum SDK C++ samples

TensorFlow implementation of "Learning from Simulated and Unsupervised Images through Adversarial Training"