Image Restoration Toolkit in PyTorch

Usage

Dependencies

PyTorch >= 1.6 and

conda install tensorboard h5py scikit-image

Train

python trainer.py --dataset [dataset&options] --eval_datasets [datasets&options] --model [model&options] --job_dir [dir]

e.g.

python trainer.py --dataset div2k --eval_datasets div2k set5 bsds100 urban100 --scale 2 --model wdsr --num_blocks 16 --job_dir ./wdsr_x2_b16

We also support multi-GPU and mixed precision training:

python -m torch.distributed.launch --nproc_per_node=[\#gpu] trainer.py --amp [options]

Option --master_port=$(($RANDOM+($RANDOM%2)*32768)) may help.

Evaluation

python trainer.py --eval_only --dataset [dataset&options] --eval_datasets [datasets&options] --model [model&options] --job_dir [dir]
# or
python trainer.py --eval_only --dataset [dataset&options] --eval_datasets [datasets&options] --model [model&options] --job_dir X --ckpt [path]

e.g.

python trainer.py --eval_only --dataset div2k --eval_datasets div2k set5 bsds100 urban100 --scale 2 --model wdsr --num_blocks 16 --job_dir ./wdsr_x2_b16
# or
python trainer.py --eval_only --dataset div2k --eval_datasets div2k set5 bsds100 urban100 --scale 2 --model wdsr --num_blocks 16 --job_dir X --ckpt ./wdsr_x2_b16/latest.pth

Datasets

DIV2K dataset: DIVerse 2K resolution high quality images as used for the NTIRE challenge on super-resolution @ CVPR 2017

Benchmarks (Set5, BSDS100, Urban100)

Download and organize data like:

pt_ir/data/DIV2K/
├── DIV2K_train_HR
├── DIV2K_train_LR_bicubic
│   └── X2
│   └── X3
│   └── X4
├── DIV2K_valid_HR
└── DIV2K_valid_LR_bicubic
    └── X2
    └── X3
    └── X4
pt_ir/data/Set5/*.png
pt_ir/data/BSDS100/*.png
pt_ir/data/Urban100/*.png

Related repos

[NSR] Neural Sparse Representation for Image Restoration

[SCN] Scale-wise Convolution for Image Restoration

[WDSR] Wide Activation for Efficient Image and Video Super-Resolution