[CVPR 2024] Boosting Spike Camera Image Reconstruction from a Perspective of Dealing with Spike Fluctuations

Rui Zhao^1,2, Ruiqin Xiong^1,2, Jing Zhao^1,2, Jian Zhang³, Xiaopeng Fan⁴, Zhaofei Yu^1,2, Tiejun Huang^1,2

1. School of Computer Science, Peking University
2. National Key Laboratory for Multimedia Information Processing, Peking University
3. School of Electronic and Computer Engineering, Peking University
4. School of Computer Science and Technology, Harbin Institute of Technology

This repository contains the official source code for our paper:

Boosting Spike Camera Image Reconstruction from a Perspective of Dealing with Spike Fluctuations

CVPR 2024

Environment

You can choose cudatoolkit version to match your server. The code is tested on PyTorch 2.0.1+cu120.

conda create -n bsf python==3.10.9
conda activate bsf
# You can choose the PyTorch version you like, we recommand version >= 1.10.1
# For example
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip3 install -r requirements.txt

Prepare the Data

1. Download the dataset (Approximate 50GB)

Link of the dataset (BaiduNetDisk) (Password: 2728)

2. Deploy the dataset for training faster (Approximate another 125GB)

firstly modify the data root and output root in ./prepare_data/crop_dataset_train.py and ./prepare_data/crop_dataset_val.py

cd prepare_data &&
bash crop_train.sh $your_gpu_id &&
bash crop_val.sh $your_gpu_id

Evaluate

CUDA_VISIBLE_DEVICES=$1 python3 -W ignore main.py \
--alpha 0.7 \
--vis-path vis/bsf \
-evp eval_vis/bsf \
--logs_file_name bsf \
--compile_model \
--test_eval \
--arch bsf \
--pretrained ckpt/bsf.pth

Train

CUDA_VISIBLE_DEVICES=$1 python3 -W ignore main.py \
-bs 8 \
-j 8 \
-lr 1e-4 \
--epochs 61 \
--train-res 96 96 \
--lr-scale-factor 0.5 \
--milestones 10 20 30 40 50 60 70 80 90 100 \
--alpha 0.7 \
--vis-path vis/bsf \
-evp eval_vis/bsf \
--logs_file_name bsf \
--compile_model \
--weight_decay 0.0 \
--eval-interval 10 \
--half_reserve 0 \
--arch bsf

Citations

If you find this code useful in your research, please consider citing our paper:

@inproceedings{zhao2024boosting,
  title={Boosting Spike Camera Image Reconstruction from a Perspective of Dealing with Spike Fluctuations},
  author={Zhao, Rui and Xiong, Ruiqin and Zhao, Jing and Zhang, Jian and Fan, Xiaopeng and Yu, Zhaofei, and Huang, Tiejun},
  booktitle={IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2024}
}