The official implementation of "Rank-One Network: An Effective Framework for Image Restoration" via TensorFlow

• Dependencies
• Quick test
  • Noise-free image super-resolution
  • Realistic image super-resolution
  • Gray-scale image denoising
  • Color image denoising
• How to train RONet
  • Pre-train RODec
  • Train RORec
• Citation

We implemented the method on Ubuntu 16.04 with Python 3.6. Before training and test, you need to create an environment via Anaconda (suppose it has been installed on your computer), and install Cuda 8.0 and CuDNN 7.1.3, as follows,

conda create -n RONet python=3.6
source activate RONet
conda install cudnn==7.1.3 # cuda 8.0 would be installed too.

Besides, you need to install the following packages using pip install -r requirements.txt

• numpy==1.15.4
• tensorflow==1.10.0
• opencv-python==4.4.0
• tqdm==4.48.2
• scikit-image==0.17.1
• Pillow==7.2.0

Finally, clone our code from the repository and start a quick test:

git clone https://github.com/shangqigao/RONet.git
cd RONet

We tested our models on the widely used benchmark datasets, you can download the collection of them from BaiduPan (extraction code: 6c1k) or OneDrive, and unzip it to the folder data/Test

Following EDSR, We test our models with four widely used bechmark datasets :

• Set5-Bevilacqua et al. BMVC 2012
• Set14-Zeyde et al. LNCS 2010
• B100-Martin et al. ICCV 2001
• Urban100-Huang et al. CVPR 2015

We provided the pre-trained model RONet-NF, please unzip it to the folder models.

Besides, we provided the script demo.sh for test, please uncomment the following code in src/demo.sh,

python RONet_test.py --dataset Set5 --input_data_dir ../data/Test/benchmark --task BiSR --upscale 4 --net_type net_sr --depth_RODec 3 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --RONet_checkpoint ../models/RONet-NF/model --save_dir ../results --GPU_ids 0

, and test RONet-NF on Set5 as follows,

cd src
sh demo.sh

Following Track2, realistic image SR of NTIRE2018, we used DIV2K_valid_LR_mild to test our models. We provided the pre-trained model RONet-R, please unzip it to the folder models.

Besides, we provided the script demo.sh for test, please uncomment the following code in src/demo.sh,

python RONet_test.py --dataset DIV2K_mild --input_data_dir ../data/Test/benchmark --task ReSR --upscale 4 --net_type net_sr --depth_RODec 3 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --RONet_checkpoint ../models/RONet-R/model --save_dir ../results --ensemble --GPU_ids 0

, and test RONet-R on DIV2K_valid_LR_mild as follow,

cd src
sh demo.sh

We tested our models with three widely used datasets:

• RNI6-Lebrun et al. IPOL 2015
• Set12
• BSD68-Roth et al. CVPR 2005

We provided the pre-trained model RONet-G_sigmaxx, where xxdenotes the noise level, please unzip it to the folder models

Besides, we provided the script demo.sh for test, please uncomment the following code in src/demo.sh,

python RONet_test.py --dataset RNI6 --input_data_dir ../data/Test/benchmark --task DEN --net_type net_den --deep_scale 48 --depth_RODec 1 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 1 --RONet_checkpoint ../models/RONet-G_sigma50/model --save_dir ../results --sigma 50 --GPU_ids 0

, and test RONet-G on RNI6 with the noise levle of 50, as follows,

cd src
sh demo.sh

We tested our models with three widely used datasets:

• CBSD68-Roth et al. CVPR 2005
• Kodak24-Franzen et al. 1999
• McMaster-Zhang et al. 2011

We provided the pre-trained model RONet-C, please unzip it to the folder models

Besides, we provided the script demo.sh for test, please uncomment the following code in src/demo.sh,

python RONet_test.py --dataset CBSD68 --input_data_dir ../data/Test/benchmark --task DEN --net_type net_den --deep_scale 16 --depth_RODec 1 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --RONet_checkpoint ../models/RONet-C/model --save_dir ../results --sigma 50 --GPU_ids 0

, and test RONet-C on CBSD68 with the noise levle of 50, as follows,

cd src
sh demo.sh

We adopted a two-step strategy to train our RONet, i.e., (1) pretrained the RO decomposition network (RODec), and (2) fixed RODec and trained the RO reconstruction network (RORec). You can download the training datasets from the official website, and put them to the folder data/Train. Besides, please download the checkpoint of VGG19 and put vgg_19.ckpt to the folder models

• Dataset: DIV2K_train_HR
• Strategy: unsupervised training
• Demo: please uncomment the following code in the provided script (src/demo.sh).

nohup python -u RODec_train.py --train_mode unsupervised --out_channel 3 --input_data_dir ../data/Train --augment --log_dir ../logs/UROD-C --GPU_ids 0 >out &

Then, please train RODec for color mode using the following code:

cd scr
sh demo.sh

Besides, we provided pre-trained RODec models for grayscale and color images:

• Noise-free image super-resolution
  • Dataset: DIV2K_train_LR_bicubic_X4
  • Strategy: supervised learning
  • Demo: please uncomment the following code in the provided script (src/demo.sh).

nohup python -u RONet_train.py --input_data_dir ../data/Train --augment --task BiSR --upscale 4 --net_type net_sr --depth_RODec 3 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --vgg_checkpoint ../models/vgg_19.ckpt --RODec_checkpoint ../models/UROD-C/model --log_dir ../logs/RONet-NF --GPU_ids 0 >out &

Then, you can train a RONet-NF model using the following code:

cd scr
sh demo.sh

You can check logging information via vim out.

• Realistic image super-resolution
  • Dataset: DIV2K_train_LR_mild
  • Strategy: supervised learning
  • Demo: please uncomment the following code in the provided script (src/demo.sh).

nohup python -u RONet_train.py --input_data_dir ../data/Train --augment --task ReSR --upscale 4 --net_type net_sr --depth_RODec 3 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --vgg_checkpoint ../models/vgg_19.ckpt --RODec_checkpoint ../models/UROD-C/model --log_dir ../logs/RONet-R --GPU_ids 0 >out &

Then, you can train a RONet-R model using the following code:

cd scr
sh demo.sh

• Gray-scale image denoising
  • Dataset: DIV2K_train_HR
  • Strategy: supervised learning
  • Demo: please uncomment the following code in the provided script (src/demo.sh).

nohup python -u RONet_train.py --input_data_dir ../data/Train --augment --task DEN --net_type net_den --deep_scale 48 --depth_RODec 1 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 1 --RODec_checkpoint ../models/UROD-G/model --sigma 50 --log_dir ../logs/RONet-G_sigma50 --GPU_ids 0 >out &

Then, you can train a RODec-G model for the noise level of 50 using the following code:

cd scr
sh demo.sh

• Color image denoising
  • Dataset: DIV2K_train_HR
  • Strategy: supervised learning
  • Demo: please uncomment the following code in the provided script (src/demo.sh).

nohup python -u RONet_train.py --input_data_dir ../data/Train --augment --task DEN --net_type net_den --deep_scale 16 --depth_RODec 1 --depth_RecROs 3 --depth_RecRes 6 --depth_RecFus 3 --out_channel 3 --sigma 75 --range --RODec_checkpoint ../models/UROD-C/model --log_dir ../models/RONet-C --GPU_ids 0 >out &

Then, you can train a RODec-C model for the noise level of 50 using the following code:

cd scr
sh demo.sh

If our work is useful in your research or publication, please cite the work:

[1] Shangqi Gao, and Xiahai Zhuang, "Rank-One Network: An Effective Framework for Image Restoration", TPAMI, 2020. [arXiv] [TPAMI]

@Article{ronet/tpami/2020,
	title={Rank-One Network: An Effective Framework for Image Restoration},
	author={S. {Gao} and X. {Zhuang}},
    journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
    doi={10.1109/TPAMI.2020.3046476},
    number=xxx,
    volumn=xxx,
    year=2020,
    pages={xx--xx}  
}

Don't hesitate to contact us via [email protected] or [email protected], if you have any questions.