PyNET-v2: A Python repository from Qirui-Y

PyNet-V2 Mobile: Efficient On-Device Photo Processing With Neural Networks

1. Overview [Paper (in progress)] [Project Webpage (in progress)]

This repository provides the implementation of further improvement of the PyNet model originally presented in this paper.

2. Prerequisites

Python: scipy, numpy, imageio and pillow packages
TensorFlow 1.X + CUDA cuDNN
Nvidia GPU

3. First steps

Download the pre-trained VGG-19 model ^Mirror and put it into vgg_pretrained/ folder.
Download the pre-trained PyNET model and put it into models/original/ folder.
Download Zurich RAW to RGB mapping dataset and extract it into raw_images/ folder.
_{This folder should contain three subfolders: train/, test/ and full_resolution/}

_{Please note that Google Drive has a quota limiting the number of downloads per day. To avoid it, you can login to your Google account and press "Add to My Drive" button instead of a direct download. Please check this issue for more information.}

4. PyNet-V2 Mobile CNN

[WIP]

5. Training the model

The model is trained level by level, starting from the lowest. The script below incorporates all training steps:

./train.sh

6. Test the provided pre-trained models on full-resolution RAW image files

python test_model_keras.py

Optional parameters:

--model: - path to the Keras model checkpoint
--inp_path: raw_images/test/ - path to the folder with Zurich RAW to RGB dataset
--out_path: . - path to the output images

7. Folder structure

models/ - logs and models that are saved during the training process
models/original/ - the folder with the provided pre-trained PyNET model
raw_images/ - the folder with Zurich RAW to RGB dataset
results/ - visual results for small image patches that are saved while training
results/full-resolution/ - visual results for full-resolution RAW image data saved during the testing
vgg-pretrained/ - the folder with the pre-trained VGG-19 network

load_dataset.py - python script that loads training data
model.py - PyNET implementation (Keras)
train_model_keras.py - implementation of the training procedure
test_model_keras.py - applying the pre-trained model to full-resolution test images
utils.py - auxiliary functions
vgg.py - loading the pre-trained vgg-19 network

9. Bonus files

These files can be useful for further experiments with the model / dataset:

dng_to_png.py - convert raw DNG camera files to PyNET's input format
ckpt2pb_keras.py - converts Keras checkpoint to TFLite format
evaluate_accuracy_tflite.py - compute PSNR and MS-SSIM scores on Zurich RAW-to-RGB dataset for TFLite model

10. License

Licensed under the CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International).

The code is released for academic research use only.

11. Citation