97 lines
3.0 KiB
Markdown
97 lines
3.0 KiB
Markdown
# RAFT
|
|
This repository contains the source code for our paper:
|
|
|
|
[RAFT: Recurrent All Pairs Field Transforms for Optical Flow](https://arxiv.org/pdf/2003.12039.pdf)<br/>
|
|
Zachary Teed and Jia Deng<br/>
|
|
|
|
<img src="RAFT.png">
|
|
|
|
## Requirements
|
|
Our code was tested using PyTorch 1.3.1 and Python 3. The following additional packages need to be installed
|
|
|
|
```Shell
|
|
pip install Pillow
|
|
pip install scipy
|
|
pip install opencv-python
|
|
```
|
|
|
|
## Demos
|
|
Pretrained models can be downloaded by running
|
|
```Shell
|
|
./scripts/download_models.sh
|
|
```
|
|
|
|
You can run the demos using one of the available models.
|
|
|
|
```Shell
|
|
python demo.py --model=models/chairs+things.pth
|
|
```
|
|
|
|
or using the small (1M parameter) model
|
|
|
|
```Shell
|
|
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](https://lmb.informatik.uni-freiburg.de/resources/datasets/FlyingChairs.en.html#flyingchairs) and [FlyingThings3D](https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html) datasets. Finetuning and evaluation require the [Sintel](http://sintel.is.tue.mpg.de/) and [KITTI](http://www.cvlibs.net/datasets/kitti/eval_scene_flow.php?benchmark=flow) datasets. We organize the directory structure as follows. By default `datasets.py` will search for the datasets in these locations
|
|
|
|
```Shell
|
|
├── 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)
|
|
|
|
```Shell
|
|
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
|
|
|
|
```Shell
|
|
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
|
|
|
|
```Shell
|
|
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
|
|
|
|
```Shell
|
|
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
|
|
|
|
```Shell
|
|
python evaluate.py --model=models/chairs+things.pth
|
|
```
|
|
|
|
or the small model by including the `small` flag
|
|
|
|
```Shell
|
|
python evaluate.py --model=models/small.pth --small
|
|
```
|