RAFT

This repository contains the source code for our paper:

RAFT: Recurrent All Pairs Field Transforms for Optical Flow
Zachary Teed and Jia Deng

Requirements

Our code was tested using PyTorch 1.3.1 and Python 3. The following additional packages need to be installed

pip install Pillow
pip install scipy
pip install opencv-python

Demos

Pretrained models can be downloaded by running

./scripts/download_models.sh

You can run the demos using one of the available models.

python demo.py --model=models/chairs+things.pth

or using the small (1M parameter) model

python demo.py --model=models/small.pth --small

Running the demos will display the two images and a vizualization of the optical flow estimate. After the images display, press any key to continue.

Training

To train RAFT, you will need to download the required datasets. The first stage of training requires the FlyingChairs and FlyingThings3D datasets. Finetuning and evaluation require the Sintel and KITTI datasets. We organize the directory structure as follows. By default datasets.py will search for the datasets in these locations

├── datasets
│   ├── Sintel
|   |   ├── test
|   |   ├── training
│   ├── KITTI
|   |   ├── testing
|   |   ├── training
|   |   ├── devkit
│   ├── FlyingChairs_release
|   |   ├── data
│   ├── FlyingThings3D
|   |   ├── frames_cleanpass
|   |   ├── frames_finalpass
|   |   ├── optical_flow

We used the following training schedule in our paper (note: we use 2 GPUs for training)

python train.py --name=chairs --image_size 368 496 --dataset=chairs --num_steps=100000 --lr=0.0002 --batch_size=6

Next, finetune on the FlyingThings dataset

python train.py --name=things --image_size 368 768 --dataset=things --num_steps=60000 --lr=0.00005 --batch_size=3 --restore_ckpt=checkpoints/chairs.pth

You can perform dataset specific finetuning

Sintel

python train.py --name=sintel_ft --image_size 368 768 --dataset=sintel --num_steps=60000 --lr=0.00005 --batch_size=4 --restore_ckpt=checkpoints/things.pth

KITTI

python train.py --name=kitti_ft --image_size 288 896 --dataset=kitti --num_steps=40000 --lr=0.0001 --batch_size=4 --restore_ckpt=checkpoints/things.pth

Evaluation

You can evaluate a model on Sintel and KITTI by running

python evaluate.py --model=models/chairs+things.pth

or the small model by including the small flag

python evaluate.py --model=models/small.pth --small