This repository contains a PyTorch reimplementation of Deep Bilateral Learning for Real-Time Image Enhancements.
This implementation supports raw images in .dng format and LDR images. A dataset folder ideally should have the following structure:
dataset
├── train
│ ├── input
│ └── output
└── eval
├── input
└── outputA small example "false color" dataset is available here.
To train a model (on GPU), run the following command:
python train.py --epochs=1000 --train_data_dir=<train_data_dir> --eval_data_dir=<eval_data_dir> --cudaUse --hdr when handling raw (.dng) images.
To test a model on a single image, run the following command:
python test.py --ckpt_path=<ckpt_path> --test_img_path=<test_img_path> --cuda- HDRnet is a lightweight model, but it requires high-resolution images for training, so data processing could be a bottleneck for speed. For generality, we did not do much optimisation in this regard, so the current data pipeline is quite slow. One possible solution is to rewrite the dataset classes to pre-load and pre-process all images before training.
- You could switch on data augmentation in
transforms.Compose, but for the same reason as above, it is very inefficient to do data augmentation on CPU. One possible solution is to do data augmentation on GPU. - Currently, the model has a fixed structure as described in the original paper. This implementation does not support a variable network structure like the official TensorFlow implementation.