Skip to content

Repository files navigation

FlowIt: Global Matching via Hierarchical Transformers and Optimal Transport for Optical Flow

Official implementation of: "FlowIt: Global Matching via Hierarchical Transformers and Optimal Transport for Optical Flow" Accepted at BMVC 2026

by Sadra Safadoust, Fabio Tosi, Matteo Poggi, and Fatma Güney

📑 Table of Contents

  1. Overview
  2. Installation
  3. Checkpoints
  4. Demo
  5. Datasets
  6. Evaluation and Submission
  7. Training
  8. Citation
  9. Acknowledgements
  10. License

🍉 Overview

We present FlowIt, a novel architecture for optical flow estimation that combines global matching with confidence and occlusion-guided refinement. At its core, FlowIt leverages a hierarchical transformer architecture that captures extensive global context, enabling the model to effectively model long-range correspondences. To overcome the limitations of localized matching, we formulate the flow initialization as an optimal transport problem. This formulation yields a highly robust initial flow field, alongside explicitly derived occlusion and confidence maps. These cues are then seamlessly integrated into a guided refinement stage, where the network actively propagates reliable motion estimates from high-confidence regions into ambiguous, low-confidence areas. FlowIt achieves state-of-the-art results on the competitive Sintel benchmark and establishes new state-of-the-art cross-dataset zero-shot generalization performance on Sintel, Spring, and LayeredFlow, while also delivering competitive performance on both the KITTI benchmark and KITTI zero-shot generalization settings.

⚒️ Installation

Create a conda environment, install pytorch compatible with your hardware and then install the requirements:

conda create -n flowit python=3.11
conda activate flowit
pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt

📦 Checkpoints

You can download our pretrained checkpoints from Google Drive.

✨ Demo

You can run our gradio demo with GUI:

python demo_gradio.py

Alternatively, you can run our demo in the terminal

python demo.py --input [INPUT_PATH] --model [ckpt] --flowit_size [MODEL_SIZE] --out_dir [OUT_DIR] --save_viz_mp4
Command Line Arguments for demo.py

--input input path. Either path to a video file or to a directory containing image files. Images will be sorted based on their names.
--model path to model checkpoint.
--flowit_size size of the model. Must be one of S, M, L, or XL.
--out_dir path to directory to save outputs.
--save_flo save output optical flow in .flo format as in Middlebury and Sintel datasets.
--save_flo5 save output optical flow in .flo5 format as in Spring dataset.
--save_kitti save output optical flow in .png format as in KITTI dataset.
--save_viz save per-frame flow visualization as .png files.
--save_viz_mp4 save flow visualizations as an mp4 video.
--fps fps of the output mp4 video, used when inputs are images. when the input is a video, its fps is used for the output.
--resize resize the input to W H (two integers, e.g. --resize 960 540) before passing into the model. The output will be scaled back to the original resolution.

🗄️ Datasets

To train and evaluate FlowIt, you need to download the required datasets: FlyingChairs, FlyingThings, Sintel, KITTI, HD1K, TartanAir, LayeredFlow, and Spring.

By default, core/datasets.py will look for these datasets in the datasets folder. You can create symbolic links to wherever you downloaded the datasets so that they are laid out as follows (see RAFT, SEA-RAFT, and FlowSeek for more details):

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

🏆 Evaluation and Submission

Assuming the pretrained models are downloaded into checkpoints/

To reproduce Table 2:

SIZE=S # change to M, L, XL for other models
python evaluate.py --dataset sintel-kitti --model checkpoints/Tartan-C-T/Tartan-C-T_FlowIt-${SIZE}.pth --flowit_size ${SIZE}

To reproduce Table 3:

SIZE=S # change to M, L, XL for other models
python evaluate.py --dataset layered-spring --model checkpoints/Tartan-C-T-TSKH/Tartan-C-T-TSKH_FlowIt-${SIZE}.pth --flowit_size ${SIZE}

To reproduce Sintel results in Table 1:

SIZE=S # change to M, L, XL for other models
python submission.py --dataset sintel --model checkpoints/Tartan-C-T-TSKH/Tartan-C-T-TSKH_FlowIt-${SIZE}.pth --flowit_size ${SIZE} --output_path sintel_submission/${SIZE}
# submit the output files to the public leaderboard

To reproduce KITTI results in Table 1:

SIZE=S # change to M, L, XL for other models
python submission.py --dataset kitti --model checkpoints/Tartan-C-T-TSKH-KITTI/Tartan-C-T-TSKH-KITTI_FlowIt-${SIZE}.pth --flowit_size ${SIZE} --output_path kitti_submission/${SIZE}
# submit the output files to the public leaderboard

🔥 Training

First download the s2m2 checkpoints and place them inside checkpoints/s2m2/

SIZE=S # change to M, L, XL for other models
# Stage 0: TartanAir pretraining
python train.py --cfg config/flowit-${SIZE}/Tartan_${SIZE}.json --expname [EXP_NAME]

# Stage 1: FlyingChairs training
python train.py --cfg config/flowit-${SIZE}/Tartan-C_${SIZE}.json --expname [EXP_NAME] --restore_ckpt [PATH_TO_CKPT_FROM_PREVIOUS_STAGE]

# Stage 2: FlyingThings training
python train.py --cfg config/flowit-${SIZE}/Tartan-C-T_${SIZE}.json --expname [EXP_NAME] --restore_ckpt [PATH_TO_CKPT_FROM_PREVIOUS_STAGE]

# Stage 3: FlyingThings + Sintel + KITTI + HD1K training
python train.py --cfg config/flowit-${SIZE}/Tartan-C-T-TSKH_${SIZE}.json --expname [EXP_NAME] --restore_ckpt [PATH_TO_CKPT_FROM_PREVIOUS_STAGE]

# Stage 4: extra KITTI finetuning
python train.py --cfg config/flowit-${SIZE}/Tartan-C-T-TSKH-KITTI_${SIZE}.json --expname [EXP_NAME] --restore_ckpt [PATH_TO_CKPT_FROM_PREVIOUS_STAGE]

The checkpoints from stage 2 are used for Sintel and KITTI generalization experiments (Table 2).
The checkpoints from stage 3 are used for Sintel submissions (Table 1) and Spring and LayeredFlow generalization experiments (Table 3).
The checkpoints from stage 4 are used for KITTI benchmark submissions (Table 1).

Command Line Arguments for train.py

arguments are read from the config file and can be overriden by the command-line arguments.

--cfg path to the config file.
--expname your desired name for the experiment which will be used for logging and saving checkpoints.
--batch_size total batch size. (NOT per-gpu).
--restore_ckpt path to the checkpoint to start the training from. If it is set to 's2m2' then the appropriate s2m2 checkpoint will be loaded from checkpoints/s2m2.
--image_size resize the image to this resolution. must be multiple of 32 for training.
--num_steps Number of training iterations.
--lr Learning rate.
--flowit_size model size. must be one of S, M, L, or XL.
--log_dir directory for logs.
--save_dir directory for saving checkpoints.
--log_freq how often to log the metrics.
--continue_step If resuming an unfinished training, specify the iteration number.
--restore_optimizer If resuming an unfinished training, specify path to the optimizer.

📜 Citation

If you find our work useful in your research, please consider citing:

@article{safadoust2026flowit,
  title={FlowIt: Global Matching via Hierarchical Transformers and Optimal Transport for Optical Flow},
  author={Safadoust, Sadra and Tosi, Fabio and Poggi, Matteo and G{\"u}ney, Fatma},
  journal={arXiv preprint arXiv:2603.28759},
  year={2026}
}

🙏 Acknowledgements

This project builds on SEA-RAFT and S2M2. We are grateful to the authors for making their code publicly available.

License

This project is released under CC BY-NC 4.0 (see LICENSE).

It builds on:

About

Source code for BMVC 2026 paper "FlowIt: Global Matching via Hierarchical Transformers and Optimal Transport for Optical Flow"

Topics

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages