图像去雨实例（一）之AttentiveGAN

 

论文阅读见：《Attentive Generative Adversarial Network for Raindrop Removal from A Single Image》

首先环境准备：

pip3 install -r requirements.txt

根据实际情况安装需要的工具即可。

Test model

In this repo I uploaded a model trained on dataset provided by the origin author origin_dataset.

The trained derain net model weights files are stored in folder model/

python3 test_model.py --weights_path model/new_model/derain_gan_2019-01-25-15-55-54.ckpt-200000 --image_path data/test_data/test_2.png

这里需要注意，直接使用源码作者提供的训练模型可能会出错，因为模型和代码不匹配。上述是我自己训练后测试的bash，仅供参考。

 

The author's results are as follows:

Test Input Image

Test Derain result image

 

My results are as follows:

Test Input Image

 

Test Derain result image

 

Test Attention Map at time 1

Test Attention Map at time 2

Test Attention Map at time 3

Test Attention Map at time 4

 

Train your own model

Data Preparation

Firstly you need to organize your training data refer to the data/training_data_example folder structure. And you need to generate a train.txt record the data used for training the model.

 

Dataset

The whole dataset can be find here:

https://drive.google.com/open?id=1e7R76s6vwUJxILOcAsthgDLPSnOrQ49K

####Training Set:

861 image pairs for training.

####Testing Set A:

For quantitative evaluation where the alignment of image pairs is good. A subset of testing set B.

####Testing Set B:

239 image pairs for testing.

 

# 将原数据集分为training ,validation  by gavin
import os
import random

import argparse

# 划分验证集训练集
_NUM_TEST = 0 #20000

parser = argparse.ArgumentParser()
parser.add_argument('--folder_path', default='/home/gavin/Dataset/attentive-gan-derainnet/train', type=str,
                    help='The folder path')
parser.add_argument('--train_filename', default='./data/training_data_example/train_test.txt', type=str,
                    help='The train filename.')
parser.add_argument('--validation_filename', default='./data/training_data_example/validation.txt', type=str,
                    help='The validation filename.')

#/home/gavin/Dataset/attentive-gan-derainnet/test_b
def _get_filenames(dataset_dir):
    photo_filenames = []
    image_list = os.listdir(dataset_dir)
    photo_filenames = [os.path.join(dataset_dir, _) for _ in image_list]
    return photo_filenames


if __name__ == "__main__":

    args = parser.parse_args()

    data_dir = os.path.join(args.folder_path,'data')

    data_dir_gt = os.path.join(args.folder_path,'gt')

    # get all file names
    photo_filenames = _get_filenames(data_dir)
    photo_file_gt = _get_filenames(data_dir_gt)

    print("size of dataset is %d" % (len(photo_filenames)))
    print("size of dataset_gt is %d" % (len(photo_file_gt)))

    photo_filenames.sort()
    photo_file_gt.sort()
    print(photo_filenames)

    training_file_names = []

    for i in range(len(photo_file_gt)):
        string_filename = photo_file_gt[i] +' ' + photo_filenames[i]
        training_file_names.append(string_filename)

    '''
    # 切分数据为测试训练集
    random.seed(0)
    random.shuffle(photo_filenames)
    training_file_names = photo_filenames[_NUM_TEST:]

    print("training file size:", len(training_file_names))
    '''

    # make output file if not existed
    if not os.path.exists(args.train_filename):
        os.mknod(args.train_filename)

    # write to file
    fo = open(args.train_filename, "w")
    fo.write("\n".join(training_file_names))
    fo.close()

    # print process
    print("Written file is: ", args.train_filename)

python3 generate_flist.py --folder_path /home/gavin/Dataset/attentive-gan-derainnet/test_b

The training samples are consist of two components. A clean image free from rain drop label image and a origin image degraded by raindrops.

All your training image will be scaled into the same scale according to the config file.

Train model

In my experiment the training epochs are 200010, batch size is 1, initialized learning rate is 0.002. About training parameters you can check the global_configuration/config.py for details.

# train
python3 train_model.py --dataset_dir data/training_data_example/

# continue train
python3 train_model.py --dataset_dir data/training_data_example/  --weights_path model/new_model/derain_gan_2019-01-25-15-55-54.ckpt-200000

# test derain_gan_2019-01-25-15-55-54.ckpt-200000
python3 test_model.py --weights_path model/new_model/derain_gan_2019-01-25-15-55-54.ckpt-200000 --image_path data/test_data/test_1.png    derain_gan_2019-01-25-15-55-54.ckpt-200000