A PyTorch implementation of SRNTT[1], which is a novel Reference-based Super-Resolution method proposed in CVPR 2019.

$ python --version
Python 3.7.4

Mainly used packages are below:

- torch
- torchvision
- kornia
- scikit-learn etc., see more in Pipfile

This Python environment is managed by pipenv, so you can create a virtual environmet like below:

$ pipenv install

$ sudo docker build -t srntt-pytorch .

$ sudo docker run --rm --runtime nvidia -it srntt-pytorch /bin/bash

$ sh scripts/train.sh

More detailed descriptions are provided in the following sections.

You can get data through the following command, which is provided by the author.

$ python download_dataset.py --dataset_name DATASE_NAME

The available dataset names are {DIV2K, CUFED, CUFED5}. DIV2K and CUFED datasets can be used as train data, and CUFED5 can be used as test data. The downloaded data will be placed in ./data directory.

In the training of SRNTT, the initial weights for content extractor is set as pre-trained SRGAN model. In this implementation, we exploit Modified-SRGAN (MSRGAN) provided by MMSR bacause it is thought as a reliable source. Let's download it as below!

$ python download_pretrained_model.py

SRNTT requires HR-LR image pairs and the corresponding reference(Ref) images, moreover the swapped feature maps culculated from LR-Ref images. The calculation is time-comsuming and has high comtational costs.

In general, the swapped feature maps of train data is culculated offline to speed up the training. The offline calculation will be done by the following command.

$ python offline_texture_swapping.py --dataroot DATSET_NAME

All of the preparation was done as above, so let's train the networks. We recommend to use --use_weights option at all times (appeared in Eq. (6)).

$ python train.py --use_weights

You can test on CUFED5 dataset as below. If you have not used --use_weights option in the training, you don't have to specify the option.

$ python test.py -w ./runs/your/pth/path --use_weights

Now, we support online inference on any images. The input image you specify will be x4 downscaled by bicubic, and super-resolved with your ref image.

$ python online_inference.py -w ./runs/your/pth/path -i ./your/input/image -r ./your/ref/image --use_weights

You can get pretrained models from here.

HR/LR image means GT/input image. MSRGAN is the result of MSRGAN from MMSR. SR image (HR) indicates super-resolved image with HR image itself, and SR image (Lx) indicates super-resolved one with a similar reference (left bottom) image.

See the other results and metrics here.

Thanks you for the following!

1. Zhang, Zhifei, et al. "Image Super-Resolution by Neural Texture Transfer." Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2019.

1. The original implementation (TF) - https://github.com/ZZUTK/SRNTT
2. MMSR - https://github.com/open-mmlab/mmsr