This is the official codebase for the paper TMA: Temporal Motion Aggregation for Event-based Optical Flow.
The DSEC dataset for optical flow can be downloaded here.
Some preprocess is helpful to save training time. We use pre-generated event volumes saved in .npz files and flows in .npy files. Basically, we follow the data preprocess in E-RAFT.
We put data in datsets/dsec_full folder, and the structure should be like this:
|-dsec_full
|-trainval
|-thun_00_a
|-000000.npz
|-flow_000000.npy
|-000001.npz
|-flow_000001.npy
...
|-zurich_city_02_a
|-000000.npz
|-flow_000000.npy
|-000001.npz
|-flow_000001.npy
...
|-zurich_city_02_d
|-test
|-interlaken_00_b
|-xxxxxx.npz
|-xxxxxx.npz
|-interlaken_01_a
|-xxxxxx.npz
|-xxxxxx.npz
For trainval data, each .npz file contains two consecutive event streams named events_prev and events_curr, each flow_xxxxxx.npy file contains corresponding 16-bit optical flow.
For test data, the .npz file is indexed by test timestamp, which is useful for generating predictions for online benchmark.
For training data, use
cd data_utils
python gen_dsec.py
You need to change event_path, flow_path and output_dir for correct data generation.
For test data, use
cd data_utils
python gen_dsec_upload.py
You need to change event_path, ts_path and output_dir for correct data generation.
sh train.sh
Please choose your expected folder name to save your checkpoints. By default, ckpts/ is used.
--checkpoint_dir : Path to save checkpoints, here we use ckpts/ for convenience.
--wandb : Optional, if you want to visualize training loss.
python train_split.py --checkpoint_dir 'your_checkpoint_dir/'
Please assign a directory to save checkpoints by --checkpoint_dir.
If you want to use wandb to visualize the loss, --wandb is optional.
We also provide a split example in datasets/dsec_split/train/split_example.txt and datasets/dsec_split/val/split_example.txt.
If you find this codebase helpful for your research, please cite our paper:
@inproceedings{liu2023tma,
title={TMA: Temporal Motion Aggregation for Event-based Optical Flow},
author={Liu, Haotian and Chen, Guang and Qu, Sanqing and Zhang, Yanping and Li, Zhijun and Knoll, Alois and Jiang, Changjun},
booktitle={ICCV},
year={2023},
}
If you have any concerns about this codebase or our paper, please feel free to drop me an E-mail.