Closing the Loop: Joint Rain Generation and Removal via Disentangled Image Translation

Torch implementation Closing the Loop: Joint Rain Generation and Removal via Disentangled Image Translation in CVPR 2021. [paper]

Demo

Rain Removal and Generation

Real Rain Removal

Real Rain Generation

Network Architecture

Architecture

Prerequisites

Linux or Windows
Python 3.6
CPU or NVIDIA GPU + CUDA CuDNN

Get Started

Installation

clone this repo
Install Pytorch 1.6.0 and other dependencies (e.g., torchvision, visdom and dominate). The requirment of main dependencies are listed in requirements.txt

Dataset and Model

Synthetic Datasets

The JRGR models for synthetic rainy cityscape dataset are here.

https://pan.baidu.com/s/1Q3nn61-LhA3Zo9eBcfNt_g

extracting code：8x9g

Real Datasets

The JRGR model along with the large number of real captured photos will be released soon.

JRGR train/test

Download our dataset or make your own dataset, the dataset should have the following structure:
```
-train
  -Ot
  -Os
  -Bs
-test
  -Ot
  -Os
  -Bs
```

The directory Ot, Os, Bs save the real rainy images, synthetic rainy images and the corresponding backgrounds of synthetic rainy images.

If you have the ground truth of real rainy images and you want to visualize them in the results, you can add Bt directory in the dataset and add the config --Bt_access 1 during training and testing.

Train the model

# Proposed training strategy: pre-train and joint train
python train.py --dataroot ./dataset/cityscape --dataset_mode rain --model unetderain --name UnetDerain
python train.py --dataroot ./dataset/cityscape --dataset_mode rain --unet_load_path ./checkpoints/UnetDerain --model raincycle --name JRGR --init_derain 1,3
# Sencondary training strategy: directly joint train
python train.py --dataroot ./dataset/cityscape --dataset_mode rain --model raincycle --name JRGR --init_derain 0

Test the model
```
python test.py --dataroot ./dataset/cityscape --dataset_mode rain --model raincycle --name JRGR
```
The test results will be saved to a html file here: ./results/RO_JRGR/latest_test/index.html.

Citation

Our code is inspired by Cycle GAN.

guyii54/JRGR