ahn1340 / jrgr

Torch implementation for CVPR 2021: Closing the Loop: Joint Rain Generation and Removal via Disentangled Image Translation.

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.