Neural Supersampling for Real-time Rendering with Pytorch

Create super-resolution images from low-resolution in real time. Non-official implementation of the paper NSRR by Facebook Reality Labs in 2020. A blog post is available with more details.

This work is based on the IMAC-projects/NSRR-PyTorch , and complete the whole function by myself, to main work I do is to organize the whole network and made some adjustments in the original code.

Environment

You need Python at least 3.5 (3.6 recommended).

To install other dependencies, you can use pip with :

pip install -r requirements.txt

Usage

Dataset

We generate the dataset by unity, if you want to get the access to our generated dataset, you can download bellow.

In order to be loaded using NSRRDataLoader, the dataset should be structured like so:

[data]
│
└───View
│   │   img_1.png
│   │   img_2.png
│    ...
│   
└───Depth
│   │   img_1.png
│   │   img_2.png
│    ...
│   
└───Motion
│   │   img_1.png
│   │   img_2.png
│    ...

Where root_dir is the data_dir in config.json of NSRRDataLoader

Note that corresponding tuples of (view, depth, motion) images files should share the same name, as they cannot be grouped together otherwise.

Train

You can remove -d 1 if you do not have a CUDA-capable GPU.

python train.py -c config.json -d 1

Specially, you can change the upsampling scale by modify the downscale_factor in config.json.

Test

Here are two pre-trained models (NSRRmodel-2 is trained only in wooden house scene, so it can only test responding scene image)：

the result will be store in the output_test folder (you can modify the storage path )

python test.py -c config.json -r /path/to/checkpoint

Result

We calculate the PSNR and SSIM of the supersampling images of different methods:

Input(180*120)	SRCNN(720*480)	NSRR (ours,720*480)

Miscellaneous information

Using :

• Pytorch project template at: https://github.com/victoresque/pytorch-template

• Pytorch implementation of SSIM: https://github.com/Po-Hsun-Su/pytorch-ssim

• Pytorch implementation of colour-space conversions: https://github.com/jorge-pessoa/pytorch-colors