Awesome

BasicSR [ESRGAN] [SFTGAN] [PPON]

:black_square_button: TODO

• Test TV loss/regularization (needs to balance loss weight with other losses).
• Test HFEN loss (needs to balance loss weight with other losses).
• Test Partial Convolution based Padding (PartialConv2D).
• Test PartialConv2D with random masks.
• Add automatic model scale change (preserve conv layers, estimate upscale layers).
• Add automatic loading of old models and new ESRGAN models.
• Downscale images before and/or after inference. Helps in cleaning up some noise or bring images back to the original scale.
• Adopt SRPGAN's extraction of features from the discriminator to test if it reduces compute usage
• Import GMFN's recurrent network and add the feature loss to their MSE model, should have better MSE results with SRGAN's features/textures (Needs testing)
• Test PPON training code. Inference is the same as the PPON repo.

Done

• [:white_check_mark:] Add on the fly augmentations (gaussian noise, blur, JPEG compression).
• [:white_check_mark:] Add TV loss/regularization options. Useful for denoising tasks, reduces Total Variation.
• [:white_check_mark:] Add HFEN loss. Useful to keep high frequency information. Used Gaussian filter to reduce the effect of noise.
• [:white_check_mark:] Add Partial Convolution based Padding (PartialConv2D). It should help prevent edge padding issues. Zero padding is the default and typically has best performance, PartialConv2D has better performance and converges faster for segmentation and classification (https://arxiv.org/pdf/1811.11718.pdf). Code has been added, but the switch makes pretained models using Conv2D incompatible. Training new models for testing. (May be able to test inpainting and denoising)
• [:white_check_mark:] Added SSIM and MS-SSIM loss functions. Originally needed to replicate the PPON training code, it can also be used on ESRGAN models
• [:white_check_mark:] Import PPON's inference network to train using BasicSR's framework. They use dilated convolutions to increase receptive field and compare against ESRGAN with perceptually good results
• [:white_check_mark:] Almost complete implementation of the PPON training, based on the original published paper. It's missing the Multiscale L1 loss in phase 2 (currently it only does the L1 calculation at full scale, together with the MS-SSIM loss). Added TV Loss to phase 1 (Content Reconstruction), HFEN to phase 2 (Structure Reconstruction) and left phase 3 (Perceptual Reconstruction) with the same GAN and VGG_Feature loss as the original and can use an alternative learning rate scheme (MultiStepLR_Restart) or the original StepLR_Restart from the paper (all these options are configurable in the JSON file). Training doesn't necessarily have to stop after finishing phase 3, but it should to be the same as in the paper.

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

1. 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.

2. 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.)

3. 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 Paper, SFTGAN repo.

4. PPON model. The model for "Progressive Perception-Oriented Network for Single Image Super-Resolution", which the authors compare favorably against ESRGAN. Training is done progressively, by freezing and unfreezing layers in phases, which are: Content Reconstruction, Structure Reconstruction and Perceptual Reconstruction. For more details, please refer to Paper, PPON repo. The pretrained model for download can also be found in the original repo.

Table of Contents

1. Dependencies
2. Codes
3. Usage
4. Datasets
5. Pretrained models

Dependencies

• Python 3 (Recommend to use Anaconda)
• PyTorch >= 0.4.0
• NVIDIA GPU + CUDA
• Python packages: pip install numpy opencv-python

Optional Dependencies

• Python package: pip install tensorboardX, for visualizing curves.
• Python package: pip install lmdb, for lmdb database support.

Codes

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

We also provide:

1. Some useful scripts. More details in ./codes/scripts.
2. Evaluation codes, e.g., PSNR/SSIM metric.
3. 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

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

Test SR models

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

Test SFTGAN models

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

Test PPON models

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

How to Train

Train ESRGAN (SRGAN) models

We use a PSNR-oriented pretrained SR model to initialize the parameters for better quality. According to the author's paper and some testing, this will also stabilize the GAN training and allows for faster convergence.

1. Prepare datasets, usually the DIV2K dataset. More details are in codes/data and (Faster IO speed).
2. Optional: If the intention is to replicate the original paper here you would prerapre the PSNR-oriented pretrained model. You can also use the original RRDB_PSNR_x4.pth as the pretrained model for that purpose, otherwise any existing model will work as pretrained.
3. Modify the configuration file options/train/train_esrgan.json
4. Run command: python train.py -opt options/train/train_esrgan.json

Train SR models

1. Prepare datasets, usually the DIV2K dataset. More details are in codes/data.
2. Modify the configuration file options/train/train_sr.json
3. 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.

1. 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.
2. Put the images and segmentation probability maps in a folder as described in codes/data.
3. 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
4. Modify the configuration file in options/train/train_sftgan.json
5. Run command: python train.py -opt options/train/train_sftgan.json

Train PPON models

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

Resuming Training

When resuming training, just pass a option with the name resume_state, like , <small>"resume_state": "../experiments/debug_001_RRDB_PSNR_x4_DIV2K/training_state/200.state". </small>

Datasets

Several common SR datasets are list below.

Any dataset can be augmented to expose the model to information that might not be available in the images, such a noise and blur. For this reason, Data Augmentation has been added to the options in this repository and it can be extended to include other types of augmentations.

Pretrained models

The most recent community pretrained models can be found in the Wiki.

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

Models that were trained using the same pretrained model or are derivates of the same pretrained model are able to be interpolated to combine the properties of both. The original author demostrated this by interpolating the PSNR pretrained model (which is not perceptually good, but results in smooth images) with the ESRGAN resulting models that have more details but sometimes is excessive to control a balance in the resulting images, instead of interpolating the resulting images from both models, giving much better results.

The authors continued exploring the capabilities of linearly interpolating models in their new work "DNI" (CVPR19): Deep Network Interpolation for Continuous Imagery Effect Transition with very interesting results and examples. The script for interpolation can be found in the net_interp.py file, but a new version with more options will be commited at a later time. This is an alternative to create new models without additional training and also to create pretrained models for easier fine tuning.

More details and explanations of interpolation can be found here in the Wiki.

Following are the original pretrained models that the authors made available for ESRGAN, SFTGAN and PPON:

For more details about the original pretrained models, please see experiments/pretrained_models.

Additional Help

If you have any questions, we have a discord server where you can ask them and a Wiki with more information.

Acknowledgement

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

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}
}
@article{Hui-PPON-2019,
    title={Progressive Perception-Oriented Network for Single Image Super-Resolution},
    author={Hui, Zheng and Li, Jie and Gao, Xinbo and Wang, Xiumei},
    booktitle={arXiv:1907.10399v1},
    year={2019}
}
@InProceedings{Liu2019abpn,
    author = {Liu, Zhi-Song and Wang, Li-Wen and Li, Chu-Tak and Siu, Wan-Chi},
    title = {Image Super-Resolution via Attention based Back Projection Networks},
    booktitle = {IEEE International Conference on Computer Vision Workshop(ICCVW)},
    month = {October},
    year = {2019}
}