yuke93/RL-Restore

RL-Restore [project page][paper]

:triangular_flag_on_post: Support arbitrary input size. Aug 25
:triangular_flag_on_post: Add Python3 compatibility. Aug 6
:triangular_flag_on_post: 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

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
:heavy_exclamation_mark: 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} 
}