/RL-Restore

Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning (CVPR 2018 Spotlight)

Primary LanguagePythonMIT LicenseMIT

RL-Restore [project page][paper]

🚩 Support arbitrary input size. Aug 25
🚩 Add Python3 compatibility. Aug 6
🚩 Training code is ready! Jun 15

Overview

  • Framework

  • Synthetic & real-world results

Prerequisite

  • Anaconda is highly recommended as you can easily adjust the environment setting.

    pip install opencv-python scipy tqdm h5py

  • We have tested our code under the following settings:

    Python TensorFlow CUDA cuDNN
    2.7 1.3 8.0 5.1
    3.5 1.4 8.0 5.1
    3.6 1.10 9.0 7.0

Test

  • Start testing on synthetic dataset

    python main.py --dataset moderate

    dataset: choose a test set among mild, moderate and severe

  • ❗ Start testing on real-world data (support arbitrary input size)

    python main.py --dataset mine
    • You may put your own test images in data/test/mine/

  • Dataset

    • All test sets can be downloaded at Google Drive or Baidu Cloud.

    • Replace test_images/ with the downloaded data and play with the whole dataset.

  • Naming rules

    • Each saved image name refers to a selected toolchain. Please refer to my second reply in this issue.

Train

  • Download training images

    • Download training images (down-sampled DIV2K images) at Google Drive or Baidu Cloud.

    • Move the downloaded file to data/train/ and unzip.

  • Generate training data

    • Run data/train/generate_train.m to generate training data in HDF5 format.

    • You may generate multiple .h5 files in data/train/

  • Let's train!

    python main.py --is_train True

    • When you observe reward_sum is increasing, it indicates training is going well.

    • You can visualize reward increasing by TensorBoard.

Acknowledgement

The DQN algorithm is modified from DQN-tensorflow.

Citation

@inproceedings{yu2018crafting,
 author = {Yu, Ke and Dong, Chao and Lin, Liang and Loy, Chen Change},
 title = {Crafting a Toolchain for Image Restoration by Deep Reinforcement Learning},
 booktitle = {Proceedings of IEEE Conference on Computer Vision and Pattern Recognition},
 pages={2443--2452},
 year = {2018} 
}