The cbunet from njuvision

cbunet's Introduction

Rendering Nighttime Image Via Cascaded Color and Brightness Compensation

RAW (visualized)

Conventional ISP

Our CBUnet

- This code is still being organized and we will try to finish it before June 26th.

1. Overview [Paper] [Project Webpage] [Dataset]

This repository provides PyTorch implementation of the Rendering Nighttime Image Via Cascaded Color and Brightness Compensation. The model is trained for rendering nighttime RAW images into RGB. More visual results could be found in our website.

2. Prerequisites

Python: requeriments.txt
PyTorch libraries
Nvidia GPU

3. Demo

Download the pre-trained denoising model dn_mwrcanet_raw_c1.pth ^*, stage_1.checkpoint and stage_2.checkpoint and put them all into .pretrain folder.
Caupture the nighttime RAW image using your phone (iPhone app: RAW+, Android app: Google Camera) or DSLR cameras.

python demo.py -i xxx.dng

4. Training the model

Download the datasets from [Dataset] and edit the NR2R_ROOT in dataset.py
ownload the pretrained stage-1 from cube++.pth and stage-2 from p4_2.pkl ^** and put them all in .pretrain folder.
The model is trained stage by stage, start training from the stage-1:

CUDA_VISIBLE_DEIVCES=0 python train_stage_1.py -b 16 -d 0

then copy the stage_1's checkpoint to .pretrain/stage_1.checkpoint

CUDA_VISIBLE_DEIVCES=0 python train_stage_2.py -b 16 -d 0

then copy the stage_2's checkpoint to .pretrain/stage_2.checkpoint

CUDA_VISIBLE_DEIVCES=0 python finetune.py -b 16 -d 0

5. Test the provided pre-trained models

python test.py --stage1 xxx.checkpoint --stage2 xxx.checkpoint --visual visualization

6. Citation

@article{ignatov2020replacing,
  title={Replacing Mobile Camera ISP with a Single Deep Learning Model},
  author={Ignatov, Andrey and Van Gool, Luc and Timofte, Radu},
  journal={arXiv preprint arXiv:2002.05509},
  year={2020}
}