Code of 'Coupled Rain Streak and Background Estimation via Separable Element-wise Attention'
YINJIE TAN, QIANG WEN, JING QIN, JIANBO JIAO, GUOQIANG HAN,SHENGFENG HE
IEEE ACCESS

Coupled Rain Streak and Background Estimation via Separable Element-wise Attention

• Ubuntu 16.04
• CUDA 10.0
• Python 3.5
• Numpy
• Pytorch 1.0
• PIL
• Matplotlib

All rain dataset can be found here. DerainZoo

If use the default options, a directory ../rain_dataset should be created to store datasets and each dataset should be placed in a sub-folder. During training, the trained dataset can be modified by option setname (corresponding to the sub-folder name) in ./options/base_options.py.

Create folders train, val and 'test' in each dataset folder. For folders train and val, input image and groundtruth image are concatenated in width dimension.

• Example:
  For 'Rain_200H' dataset, if we name its sub-folder '200h', the path for training data and validation data is ../rain_dataset/200h/train and ../rain_dataset/200h/val

• Train a model:

python3 ./train.py 

• Test a model:

python3 ./test.py 

More options can be modified at directory ./options.<br/ The detail code for our model can be found at ./models/networks.py and seaunet_model.py. The training model and log of losses are stored in ./checkpoints and the validation results (PSNR,SSIM) are stored in ./results.

The framework of this code is based on [pytorch-CycleGAN-and-pix2pix] (https://github.com/amikey/pytorch-CycleGAN-and-pix2pix).