/SEA-RAFT

Primary LanguagePythonBSD 3-Clause "New" or "Revised" LicenseBSD-3-Clause

SEA-RAFT

We introduce SEA-RAFT, a more simple, efficient, and accurate RAFT for optical flow. Compared with RAFT, SEA-RAFT is trained with a new loss (mixture of Laplace). It directly regresses an initial flow for faster convergence in iterative refinements and introduces rigid-motion pre-training to improve generalization. SEA-RAFT achieves state-of-the-art accuracy on the Spring benchmark with a 3.69 endpoint-error (EPE) and a 0.36 1-pixel outlier rate (1px), representing 22.9% and 17.8% error reduction from best published results. In addition, SEA-RAFT obtains the best cross-dataset generalization on KITTI and Spring. With its high efficiency, SEA-RAFT operates at least 2.3x faster than existing methods while maintaining competitive performance.

If you find SEA-RAFT useful for your work, please consider citing our academic paper:

Yihan Wang, Lahav Lipson, Jia Deng

@article{wang2024sea,
  title={SEA-RAFT: Simple, Efficient, Accurate RAFT for Optical Flow},
  author={Wang, Yihan and Lipson, Lahav and Deng, Jia},
  journal={arXiv preprint arXiv:2405.14793},
  year={2024}
}

Requirements

Our code is developed with pytorch 2.2.0, CUDA 12.2 and python 3.10.

conda create --name SEA-RAFT python=3.10.13
conda activate SEA-RAFT
pip install -r requirements.txt

Model Zoo

Please download the models from google drive and put them into the models folder.

Required Data

To evaluate/train SEA-RAFT, you will need to download the required datasets.

By default datasets.py will search for the datasets in these locations. You can create symbolic links to wherever the datasets were downloaded in the datasets folder. Please check RAFT for more details.

├── datasets
    ├── Sintel
    ├── KITTI
    ├── FlyingChairs/FlyingChairs_release
    ├── FlyingThings3D
    ├── HD1K
    ├── spring
        ├── test
        ├── train
        ├── val
    ├── tartanair

Training, Evaluation, and Submission

Please refer to scripts/train.sh, scripts/eval.sh, and scripts/submission.sh for more details.

Acknowledgements

This project would not have been possible without relying on some awesome repos: RAFT, unimatch, Flowformer, ptlflow, and LoFTR. We thank the original authors for their excellent work.