gmalivenko / PyNET-v2

Generating RGB photos from RAW image files with PyNET-v2 Mobile

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

Copyright (C) 2022 Andrey Ignatov. All rights reserved.

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

[WIP]

12. Any further questions?

Please contact Andrey Ignatov ([e-mail]) for more information