LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation, CVPR 2018 (Spotlight)
Switch branches/tags
Nothing to show
Clone or download

README.md

LiteFlowNet

This repository (https://github.com/twhui/LiteFlowNet) is the offical release of LiteFlowNet for my paper LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation in CVPR18 (Spotlight). The up-to-date version of the paper is available on arXiv.

It comes as the modified Caffe from DispFlowNet and FlowNet2 with our new layers, scripts, and trained models.

For more details about LiteFlowNet, please visit my project page.

License and Citation 

All code and other materials (including but not limited to the paper, figures, and tables) are provided for research purposes only and without any warranty. Any commercial use requires our consent. When using any parts of the code package or the paper (LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation) in your work, please cite the following paper:

@InProceedings{hui18liteflownet,    
 author = {Tak-Wai Hui and Xiaoou Tang and Chen Change Loy},    
 title = {LiteFlowNet: A Lightweight Convolutional Neural Network for Optical Flow Estimation},    
 booktitle  = {Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},    
 year = {2018},    
 url = {http://mmlab.ie.cuhk.edu.hk/projects/LiteFlowNet/} 
}

Prerequisites

Installation was tested under Ubuntu 14.04.5/16.04.2 with CUDA 8.0, cuDNN 5.1 and openCV 2.4.8/3.1.0.

Edit Makefile.config (and Makefile) if necessary in order to fit your machine's settings.

For openCV 3+, you may need to change opencv2/gpu/gpu.hpp to opencv2/cudaarithm.hpp in /src/caffe/layers/resample_layer.cu.

If your machine installed a newer version of cuDNN, you do not need to downgrade it. You can do the following trick:

  1. Download cudnn-8.0-linux-x64-v5.1.tgz and untar it to a temp folder, say cuda-8-cudnn-5.1

  2. Rename cudnn.h to cudnn-5.1.h in the folder /cuda-8-cudnn-5.1/include

  3. $ sudo cp cuda-8-cudnn-5.1/include/cudnn-5.1.h /usr/local/cuda/include/	
    $ sudo cp cuda-8-cudnn-5.1/lib64/lib* /usr/local/cuda/lib64/
  4. Replace #include <cudnn.h> to #include <cudnn-5.1.h> in /include/caffe/util/cudnn.hpp.

Compiling

$ cd LiteFlowNet
$ make -j 8 tools pycaffe

Datasets

  1. FlyingChairs dataset (31GB) and train-validation split.
  2. RGB image pairs (clean pass) (37GB) and flow fields (311GB) for Things3D dataset.
  3. Sintel dataset (clean + final passes) (5.3GB).
  4. KITTI12 dataset (2GB) and KITTI15 dataset (2GB) (Simple registration is required).
FlyingChairs FlyingThings3D Sintel KITTI
Crop size 448 x 320 768 x 384 768 x 384 896 x 320
Batch size 8 4 4 4

Feature warping (f-warp) layer

The source files include /src/caffe/layers/warp_layer.cpp, /src/caffe/layers/warp_layer.cu, and /include/caffe/layers/warp_layer.hpp.

The grid pattern that is used by f-warp layer is generated by a grid layer. The source files include /src/caffe/layers/grid_layer.cpp and /include/caffe/layers/grid_layer.hpp.

Feature-driven local convolution (f-lcon) layer

It is implemented using off-the-shelf components. More details can be found in /models/testing/depoly.prototxt or /models/training_template/train.prototxt.template by locating the code segment NetE-R.

Other layers

Two custom layers (ExpMax and NegSquare) are optimized in speed for forward-pass.

Training

  1. Prepare the training set. In /data/make-lmdbs-train.sh, change YOUR_TRAINING_SET and YOUR_TESTING_SET to your favourite dataset.
$ cd LiteFlowNet/data
$ ./make-lmdbs-train.sh
  1. Copy files from /models/training_template to a new model folder (e.g. NEW). Edit all the files and make sure the settings are correct for your application.
$ mkdir LiteFlowNet/models/NEW
$ cd LiteFlowNet/models/NEW
$ cp ../training_template/solver.prototxt.template solver.prototxt	
$ cp ../training_template/train.prototxt.template train.prototxt
$ cp ../training_template/train.py.template train.py
  1. Create a soft link in your new model folder
$ ln -s ../../build/tools bin
  1. Run the training script
$ ./train.py -gpu 0 2>&1 | tee ./log.txt

Trained models

The trained models (liteflownet, liteflownet-ft-sintel, liteflownet-ft-kitti) are available in the folder /models/trained. Untar the files to the same folder before you use it.

Testing

  1. Open the testing folder
$ cd LiteFlowNet/models/testing
  1. Create a soft link in the folder /testing
$ ln -s ../../build/tools bin
  1. Replace MODE in ./test_MODE.py to batch if all the images has the same resolution (e.g. Sintel dataset), otherwise replace it to iter (e.g. KITTI dataset).

  2. Replace MODEL in line 10 (cnn_model = 'MODEL') of test_MODE.py to one of the trained models (e.g. liteflownet-ft-sintel).

  3. Run the testing script. Flow fields (MODEL-0000000.flo, MODEL-0000001.flo, ... etc) are stored in the folder /testing/results having the same order as the image pair sequence.

$ test_MODE.py img1_pathList.txt img2_pathList.txt results

Evaluation

  1. End-point error (EPE) per image can be calculated using the provided script /models/testing/util/endPointErr.m

  2. Average end-point error (AEE) is simply computed by taking the average of all EPE.