Official Implementation of Compatible Training Objective for Improving Perceptual Super-Resolution in Tensorflow 2.0+.

This repository contains the implementation and training of the methods proposed in the paper Compatible Training Objective for Improving Perceptual Super-Resolution.(Link)

The methods presented in our paper were implemented with the ESRGAN network from ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks by Xintao Wang et al. In our work we propose the following:

• We provide weigthed random noise to the generator to provide it with the ability to generate diverse outputs.
• We propose a weaker content loss that is compatible with the multiple outputs of the generator, and does not contradict the adversarial loss.
• We improve the SR quality by filtering blurry regions in the training data using Laplacian activation.
• We additionally provide the LR image to the discriminator as a reference image to give better gradient feedback to the generator.

To install requirements:

pip install -r requirements.txt

You can modify the configurations of our models in ./configs/*.yaml for training and testing, which like below.

# general setting
batch_size: 16
input_size: 32
gt_size: 128
ch_size: 3
scale: 4
log_dir: '/content/drive/MyDrive/ESRGAN'
pretrain_dir: '/content/drive/MyDrive/ESRGAN-MSE-500K-DIV2K'  # directory to load from at initial training
cycle_mse: True
# generator setting
network_G:
    name: 'RRDB'
    nf: 64
    nb: 23
    apply_noise: True
# discriminator setting
network_D:
    nf: 64

# dataset setting
train_dataset:
    path: '/content/drive/MyDrive/data/div2k_hr/DIV2K_train_HR'
    num_samples: 32208
    using_bin: True
    using_flip: True
    using_rot: True
    detect_blur: True
    buffer_size: 1024           # max size of buffer
    patch_per_image: 128        # number of patches to extract from each image
test_dataset:
    set5: './test data/Set5'
    set14: './test data/Set14'

# training setting
niter: 400000

lr_G: !!float 1e-4
lr_D: !!float 1e-4
lr_steps: [50000, 100000, 200000, 300000]
lr_rate: 0.5

adam_beta1_G: 0.9
adam_beta2_G: 0.99
adam_beta1_D: 0.9
adam_beta2_D: 0.99

w_pixel: !!float 1e-2
pixel_criterion: l1

w_feature: 1.0
feature_criterion: l1

w_gan: !!float 5e-3
gan_type: ragan  # gan | ragan
refgan: True       # provide reference image

save_steps: 5000

# logging settings
logging:
    psnr: True
    lpips: True
    ssim: True
    plot_samples: True

cycle_mse: use cycle-consistent content loss

network_G/apply_noise: provide random noise to the generator network

train_dataset/detect_blur: filter blurry images in the training dataset

refgan: provide reference image to the discriminator network

Explanation of config files:

• esrgan.yaml: baseline ESRGAN (configuration(c))
• esrrefgan.yaml: +refgan (configuration(d))
• use_noise.yaml: +use noise (configuration(e))
• cyclegan.yaml: +cycle loss (configuration(f))
• cyclegan_only.yaml: -perceptual loss (configuration(g))

The training process is divided into two parts; pretraining the model with pixel-wise loss, and training the pretrained PSNR model with ESRGAN loss.

Pretrain the PSNR RDDB model.

python train_psnr.py --cfg_path="./configs/psnr.yaml" --gpu=0

Train the ESRGAN model with the pretrain PSNR model.

python train_esrgan.py --cfg_path="./configs/esrgan.yaml" --gpu=0

Configure the dataset directory and log directory in the config file before training. The DIV2K dataset is available here and the DIV8K dataset is avilable here.

python test.py --model=weights/ESRGAN-cyclemixing --gpu=0 --img_path=photo/baby.png --down_up=True --scale=4(optional)

When the down_up option is True, the image will be arbitrarily downsampled and processed through the network. For real use cases, the option must be marked False for the model to upsample the image.

All our trained models and tensorboard logs in the experiment can be downloaded here. Three trained models are included in the repository in ./weights/*.

Our methods were evaluated on LPIPS, PSNR, and SSIM using the Set5, Set14, BSD100, Urban100, and Manga109 dataset. The scores are displayed in the tables below, in the order LPIPS/PSNR/SSIM.