[AAAI 2024] Optical Flow for Spike Camera with Hierarchical Spatial-Temporal Spike Fusion

Rui Zhao¹, Ruiqin Xiong¹, Jian Zhang², Xinfeng Zhang³, Zhaofei Yu^1,4, Tiejun Huang^1,4

1. National Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University
2. School of Electronic and Computer Engineering, Peking University Shenzhen Graduate School
3. School of Computer Science and Technology, University of Chinese Academy of Sciences
4. Institute for Artificial Intelligence, Peking University

This repository contains the official source code for our paper:

Optical Flow for Spike Camera with Hierarchical Spatial-Temporal Spike Fusion

AAAI 2024

Environment

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

conda create -n hist python==3.10.9
conda activate hist
# 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 and deploy the SPIFT and PHM dataset

Link of SPIFT and PHM (BaiduNetDisk) (password: 5331)

2. Set the path of SPIFT and PHM dataset in your server

In lines 2-3 of configs/flow.yml

3. Pre-processing for SPIFT and PHM

cd encoding &&
python encoding_dat.py --data_root 'your_root' --dataset spift --data_length 25 &&
python encoding_dat.py --data_root 'your_root' --dataset phm --data_length 25

Evaluate

python main.py --arch hist_sflow --eval --pretrained ckpt/hist_sflow.pth

You can also inference optical flow from real data through demo.py

Train

python3 main.py \ 
--learning_rate 1e-4 \ 
--configs ./configs/flow.yml \
--arch hist_sflow \
--decay_factor 0.8 \
--vis_path ./vis/hist_sflow \
--save_name hist_sflow \
--eval_vis ./eval_vis/hist_sflow \
--weight_rec_loss 0.5 \
--scene_weight_list_type 1

We recommended to redirect the output logs by adding >> hist_sflow.txt 2>&1 to the last of the above command for management.

Citations

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

@inproceedings{zhao2024optical,
  title={Optical Flow for Spike Camera with Hierarchical Spatial-Temporal Spike Fusion},
  author={Zhao, Rui and Xiong, Ruiqin and Zhang, Jian and Zhang, Xinfeng and Yu, Zhaofei and Huang, Tiejun},
  booktitle={AAAI Conference on Artificial Intelligence (AAAI)},
  year={2024}
}

Acknowledgement

Parts of this code were derived from askerlee/craft. Please also consider to cite CRAFT if you'd like to cite our paper.