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:

SEA-RAFT: Simple, Efficient, Accurate RAFT for Optical Flow

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.

Custom Usage

We provide an example in custom.py. By default, this file will take two RGB images as the input and provide visualizations of the optical flow and the uncertainty.

python custom.py --cfg config/eval/spring-M.json --model models/Tartan-C-T-TSKH-spring540x960-M.pth

Datasets

To evaluate/train SEA-RAFT, you will need to download the required datasets: FlyingChairs, FlyingThings3D, Sintel, KITTI, HD1K, TartanAir, and Spring.

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 relies on code from existing repositories: RAFT, unimatch, Flowformer, ptlflow, and LoFTR. We thank the original authors for their excellent work.

wanyc-rt/SEA-RAFT