WangJunjie97 / CBUnet

Official code of the "Rendering Nighttime Image Via Cascaded Color and Brightness Compensation"

Rendering Nighttime Image Via Cascaded Color and Brightness Compensation

RAW (visualized)

Conventional ISP

Our CBUnet

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.pth

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

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

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}
}

7. Contact

Zhihao Li (lizhihao6_at_outlook.com)

8. Thaks

Zhihao Li (lizhihao6_at_outlook.com)