/SplatFlow

[IJCV 2024] SplatFlow: Learning Multi-frame Optical Flow via Splatting

Primary LanguagePythonMIT LicenseMIT

SplatFlow: Learning Multi-frame Optical Flow via Splatting

This repository contains the source code for our paper:

  • SplatFlow: Learning Multi-frame Optical Flow via Splatting (IJCV 2024) | Paper
    • We propose a novel MOFE framework SplatFlow designed explicitly for the single-resolution iterative two-frame backbones.
    • Compared with the original backbone, SplatFlow has significantly higher estimation accuracy, especially in occluded regions, while maintaining a high inference speed.
    • At the time of submission, our SplatFlow achieved state-of-the-art results on both the Sintel and KITTI2015 benchmarks, especially with surprisingly significant 19.4% error reductions compared to the previous best result submitted on the Sintel benchmark.

Updates

  • [2024.04.24] 📣 The code of SplatFlow is now available!
  • [2024.01.02] 📣 The paper of SplatFlow is accepted by IJCV 2024!

Environment

Our code has been successfully tested in the following environments:

  • NVIDIA 3090 GPU
  • CUDA 11.1
  • Python 3.8
  • PyTorch 1.8.2
conda create -n splatflow python=3.8
conda activate splatflow

pip install torch==1.8.2 torchvision==0.9.2 --extra-index-url https://download.pytorch.org/whl/lts/1.8/cu111
pip install einops==0.4.1
pip install cupy-cuda111
pip install pillow==9.5.0
pip install opencv-python==4.1.2.30

Trained Weights

Download the weights below and put them in the exp/0-pretrain path.

Model Training process Weights Comments
SplatFlow K-finetune splatflow_kitti_50k.pth
Huggingface & BaiduNetdisk
Best performance on KITTI

Demo

  • Quick start.
    bash script/demo.sh

Datasets

To train / test SplatFlow, you will need to download the required datasets and update data_root in data/dataset.py.

data_root/
│
├─ FlyingThings3D/
│  ├─ frames_cleanpass/
│  ├─ frames_finalpass/
│  └─ optical_flow/
│
├─ Sintel/
│  ├─ training/
│  └─ test/
│
├─ KITTI/
│  ├─ training/
│  └─ testing/
│
├─ HD1k/
│  ├─ hd1k_input/
│  └─ hd1k_flow_gt/
│
└─ demo/
   ├─ image/
   └─ pred/

Training

  • Train SplatFlow under the C+T training process.

    bash script/train_things.sh
  • Train SplatFlow under the S-finetune training process.

    bash script/train_sintel.sh
  • Train SplatFlow under the K-finetune training process.

    bash script/train_kitti.sh

Testing

  • Test SplatFlow on Things.

    bash script/test_things.sh
  • Test SplatFlow on KITTI.

    bash script/test_kitti.sh

Acknowledgments

We would like to thank RAFT, GMA and SoftSplat for publicly releasing their code and data.

Citing this Work

If you find our repository useful, please consider giving it a star ⭐ and citing our paper in your work:

@article{wang2024splatflow,
  title={SplatFlow: Learning Multi-frame Optical Flow via Splatting},
  author={Wang, Bo and Zhang, Yifan and Li, Jian and Yu, Yang and Sun, Zhenping and Liu, Li and Hu, Dewen},
  journal={International Journal of Computer Vision},
  pages={1--23},
  year={2024},
  publisher={Springer}
}