BasicSR [ESRGAN] [SFTGAN]

An image super-resolution toolkit flexible for development. It now provides:

PSNR-oriented SR models (e.g., SRCNN, SRResNet and etc). You can try different architectures, e.g, ResNet Block, ResNeXt Block, Dense Block, Residual Dense Block, Poly Block, Dual Path Block, Squeeze-and-Excitation Block, Residual-in-Residual Dense Block and etc.

Enhanced SRGAN model (It can also train the SRGAN model). Enhanced SRGAN achieves consistently better visual quality with more realistic and natural textures than SRGAN and won the first place in the PIRM2018-SR Challenge. For more details, please refer to Paper, ESRGAN repo. (If you just want to test the model, ESRGAN repo provides simpler testing codes.)

SFTGAN model. It adopts Spatial Feature Transform (SFT) to effectively incorporate other conditions/priors, like semantic prior for image SR, representing by segmentation probability maps. For more details, please refer to Papaer, SFTGAN repo.

BibTex

@InProceedings{wang2018esrgan,
    author = {Wang, Xintao and Yu, Ke and Wu, Shixiang and Gu, Jinjin and Liu, Yihao and Dong, Chao and Qiao, Yu and Loy, Chen Change},
    title = {ESRGAN: Enhanced super-resolution generative adversarial networks},
    booktitle = {The European Conference on Computer Vision Workshops (ECCVW)},
    month = {September},
    year = {2018}
}
@InProceedings{wang2018sftgan,
    author = {Wang, Xintao and Yu, Ke and Dong, Chao and Loy, Chen Change},
    title = {Recovering realistic texture in image super-resolution by deep spatial feature transform},
    booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    month = {June},
    year = {2018}
}

Dependencies
Codes
Usage
Datasets
Pretrained models

Dependencies

Python 3 (Recommend to use Anaconda)
PyTorch >= 0.4.0
NVIDIA GPU + CUDA
Python packages: pip install numpy opencv-python lmdb
[option] Python packages: pip install tensorflow tensorboard_logger, for visualizing curves.

Codes

./codes. We provide a detailed explaination of the code framework in ./codes.

We also provides:

Some useful scripts. More details in ./codes/scripts.
Evaluation codes, e.g., PSNR/SSIM metric.
Wiki, e.g., How to make high quality gif with full (true) color, Matlab bicubic imresize and etc.

Usage

Data and model preparation

The common SR datasets can be found in Datasets. Detailed data preparation can be seen in codes/data.

We provide pretrained models in Pretrained models.

How to Test

Test ESRGAN (SRGAN) models

Modify the configuration file options/test/test_esrgan.json
Run command: python test.py -opt options/test/test_esrgan.json

Test SR models

Modify the configuration file options/test/test_sr.json
Run command: python test.py -opt options/test/test_sr.json

Test SFTGAN models

Obtain the segmentation probability maps: python test_seg.py
Run command: python test_sftgan.py

How to Train

Train ESRGAN (SRGAN) models

We use a PSNR-oriented pretrained SR model to initialize the parameters for better quality.

Prepare datasets, usually the DIV2K dataset. More details are in codes/data and wiki (Faster IO speed).
Prerapre the PSNR-oriented pretrained model. You can use the RRDB_PSNR_x4.pth as the pretrained model.
Modify the configuration file options/train/train_esrgan.json
Run command: python train.py -opt options/train/train_esrgan.json

Train SR models

Prepare datasets, usually the DIV2K dataset. More details are in codes/data.
Modify the configuration file options/train/train_sr.json
Run command: python train.py -opt options/train/train_sr.json

Train SFTGAN models

Pretraining is also important. We use a PSNR-oriented pretrained SR model (trained on DIV2K) to initialize the SFTGAN model.

First prepare the segmentation probability maps for training data: run test_seg.py. We provide a pretrained segmentation model for 7 outdoor categories in Pretrained models. We use Xiaoxiao Li's codes to train our segmentation model and transfer it to a PyTorch model.
Put the images and segmentation probability maps in a folder as described in codes/data.
Transfer the pretrained model parameters to the SFTGAN model.
1. First train with debug mode and obtain a saved model.
2. Run transfer_params_sft.py to initialize the model.
3. We provide an initialized model named sft_net_ini.pth in Pretrained models
Modify the configuration file in options/train/train_sftgan.json
Run command: python train.py -opt options/train/train_sftgan.json

Datasets

Several common SR datasets are list below.

Name	Datasets	Short Description	Google Drive	Baidu Drive
Classical SR Training	T91	_{91 images for training}	Google Drive	Baidu Drive
	BSDS200	_{A subset (train) of BSD500 for training}
	General100	_{100 images for training}
Classical SR Testing	Set5	_{Set5 test dataset}
	Set14	_{Set14 test dataset}
	BSDS100	_{A subset (test) of BSD500 for testing}
	urban100	_{100 building images for testing (regular structures)}
	manga109	_{109 images of Japanese manga for testing}
	historical	_{10 gray LR images without the ground-truth}
2K Resolution	DIV2K	_{proposed in NTIRE17(800 train and 100 validation)}	Google Drive	Baidu Drive
	Flickr2K	_{2650 2K images from Flickr for training}
	DF2K	_{A merged training dataset of DIV2K and Flickr2K}
OST (Outdoor Scenes)	OST Training	_{7 categories images with rich textures}	Google Drive	Baidu Drive
OST (Outdoor Scenes)	OST300	_{300 test images of outdoor scences}	Google Drive	Baidu Drive
PIRM	PIRM	_{PIRM self-val, val, test datasets}	Google Drive	Baidu Drive

Pretrained models

We provide some pretrained models. More details about the pretrained models, please see experiments/pretrained_models.

You can put the downloaded models in the experiments/pretrained_models folder.

Name	Modeds	Short Description	Google Drive	Baidu Drive
ESRGAN	RRDB_ESRGAN_x4.pth	_{final ESRGAN model we used in our paper}	Google Drive	Baidu Drive
ESRGAN	RRDB_PSNR_x4.pth	_{model with high PSNR performance}	Google Drive	Baidu Drive
SFTGAN	segmentation_OST_bic.pth	_{segmentation model}	Google Drive	Baidu Drive
	sft_net_ini.pth	_{sft_net for initilization}
	sft_net_torch.pth	_{SFTGAN Torch version (paper)}
	SFTGAN_bicx4_noBN_OST_bg.pth	_{SFTGAN PyTorch version}
SRGAN^*1	SRGAN_bicx4_303_505.pth	_{SRGAN(with modification)}	Google Drive	Baidu Drive
SRResNet^*2	SRGAN_bicx4_303_505.pth	_{SRGAN(with modification)}	Google Drive	Baidu Drive

😆 Image Viewer - HandyViewer

May try HandyViewer - an image viewer that you can switch image with a fixed zoom ratio, easy for comparing image details.

Acknowledgement

Code architecture is inspired by pytorch-cyclegan.
Thanks to Wai Ho Kwok, who contributes to the initial version.

sachit-menon / BasicSR