zhanggengchen / BRVE

[CVPR 2024] Binarized Low-light Raw Video Enhancement
MIT License
23 stars 3 forks source link

Binarized Low-light Raw Video Enhancement

This repository contains the official implementation of the following paper:

Gengchen Zhang, Yulun Zhang, Xin Yuan and Ying Fu. "Binarized Low-light Raw Video Enhancement", In CVPR, 2024.

🎥 Demo

📖 Introduction

In this work, we design a binary raw video enhancement model (BRVE). It can achieve satisfactory results with low memory and computation cost.

Our BRVE has two main innovations:

🛠️ Installation

Create conda pytorch environment

conda create -n BRVE python=3.7
conda activate BRVE
conda install pytorch=1.11 torchvision cudatoolkit=11.3 -c pytorch

Install mmcv

pip install openmim
mim install mmcv-full==1.6.0

Setup mmedit with BRVE

git clone https://github.com/zhanggengchen/BRVE.git
cd BRVE
pip install -v -e .

Install dependencies

pip install -r requirements/brve.txt

💻 Usage

Prepare Datasets

Download datasets

Create the folder to place datasets.

mkdir datasets

The LLRVD dataset can be downloaded from Baidu Disk (code: xydx).

The SMOID dataset is currently not available in this repository.

Directly download datasets to datasets/ or use symbolic link.

ln -s path_to_LLRVD datasets/LLRVD
ln -s path_to_SMOID datasets/SMOID
Datasets should be organized as follows (Click to expand) : ``` |--datasets | |--LLRVD | | |--ISP_info | | | |--1 | | | | |--1.npy | | | | |--2.npy | | | | ... | | | |--2 | | | | |--1.npy | | | | |--2.npy | | | | ... | | | ... | | |--ratio_gt | | | |--1 | | | | |--gt.npy | | | |--2 | | | | |--gt.npy | | | ... | | |--ratio_noisy_all | | | |--1 | | | | |--125 | | | | | |--lq.npy | | | | |--160 | | | | | |--lq.npy | | | | |--250 | | | | | |--lq.npy | | | |--2 | | | | |--100 | | | | | |--lq.npy | | | | |--125 | | | | | |--lq.npy | | | | |--200 | | | | | |--lq.npy | | | ... | |--SMOID | | |--_Scene1 | | | |--Gain0Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | | |--Gain15Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | | |--Gain30Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | |--_Scene2 | | | |--Gain0Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | | |--Gain15Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | | |--Gain30Gain0 | | | | |--Bright | | | | | |--RAW | | | | | | |--raw_target.npy | | | | |--Dark | | | | | |--RAW | | | | | | |--raw_input.npy | | ... ```

Test

Prepare pretrained models

Create the folder to place pretrained models.

mkdir pretrained_models

The BRVE model trained on LLRVD dataset can be downloaded from Google Drive or Baidu Disk (code: 4v88).

The BRVE model trained on SMOID dataset can be downloaded from Google Drive or Baidu Disk (code: wq1u).

Put the downloaded pretrained models to pretrained_models/

Test on the LLRVD dataset

Run the following command to reproduce the results in Table 1.

torchrun --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --nproc_per_node=1 --master_port=29500 tools/test.py configs/BRVE_LLRVD.py pretrained_models/BRVE_LLRVD.pth --seed 0 --out test_results/BRVE_LLRVD/result.json --save-path test_results/BRVE_LLRVD --launcher pytorch

Notes

Set --nproc_per_node=2 to use 2 GPUs for distributed test.

The path to configuration file configs/BRVE_LLRVD.py is required.

The path to pretrained model pretrained_models/BRVE_LLRVD.pth is required.

The path to save the quantitative results is specified by --out test_results/BRVE_LLRVD/result.json.

The path to save enhanced frames and videos is specified by --save-path test_results/BRVE_LLRVD. The argument --save-path is optional.

Test on the SMOID dataset

Run the following command to reproduce the results in Table 1.

torchrun --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --nproc_per_node=1 --master_port=29500 tools/test.py configs/BRVE_SMOID.py pretrained_models/BRVE_SMOID.pth --seed 0 --out test_results/BRVE_SMOID/result.json --save-path test_results/BRVE_SMOID --launcher pytorch

Notes

Set --master_port=29501 if there is a test process running on port 29500.

Train

Train on the LLRVD dataset

torchrun --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --nproc_per_node=1 --master_port=29500 tools/train.py configs/BRVE_LLRVD.py --seed 0 --launcher pytorch

Train on the SMOID dataset

torchrun --nnodes=1 --node_rank=0 --master_addr="127.0.0.1" --nproc_per_node=2 --master_port=29500 tools/train.py configs/BRVE_SMOID.py --seed 0 --launcher pytorch

Notes

We use 2 GPUs to train BRVE model on SMOID dataset.

Evaluating the complexity of BRVE

We input a video with 100 frames and a resolution of $256 \times 256$ (in Bayer pattern, $4 \times 256 \times 256$ after packing) to evaluate the complexity of BRVE.

python tools/get_flops_bnn.py configs/BRVE_LLRVD.py --shape 100 4 128 128

❤️ Acknowledgement

Our code is developed based on BasicVSR++, Shift-Net and MMEditing (MMagic now).

The implementation of BNN is based on BBCU and IRNet.

We thank these public repositories for their excellent code.

🖊️ Citation

@inproceedings{BRVE_2024_CVPR,
  title={Binarized Low-light Raw Video Enhancement},
  author={Zhang, Gengchen and Zhang, Yulun and Yuan, Xin and Fu, Ying},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={25753-25762},
  year={2024}
}