FLOT: Scene Flow on Point Clouds guided by Optimal Transport
Gilles Puy1,
Alexandre Boulch1,
Renaud Marlet1,2
1valeo.ai, France and 2ENPC, France
If you find this code or work useful, please cite our paper:
@inproceedings{puy20flot,
title={{FLOT}: {S}cene {F}low on {P}oint {C}louds {G}uided by {O}ptimal {T}ransport},
author={Puy, Gilles and Boulch, Alexandre and Marlet, Renaud},
booktitle={European Conference on Computer Vision}
year={2020}
}
- python = 3.7
- pytorch = 1.3.1
- cuda = 10.0
- tqdm
- tensorboard
You can create a conda environment for FLOT by typing:
$ conda create --name FLOT
$ conda activate FLOT
$ conda install python=3.7 tqdm tensorboard
$ conda install pytorch=1.3.1 cuda100 torchvision -c pytorch
- Clone this repository:
$ git clone https://github.com/valeoai/FLOT.git
- Install the repository:
$ pip install -e ./FLOT
You can edit flot's code on the fly and import function and classes of flot in other project as well.
- To uninstall this package, run:
$ pip uninstall flot
By default, the datasets are stored in /path/to/flot/data
. The datasets are prepared using FlyingThing3D and the KITTI scene flow dataset. There are two different ways to prepare the datasets.
- FlowNet3d's version: Please download the datasets available here to obtain the FlyingThings3D and KITTI datasets used in this paper. These datasets are denoted FT3Do and KITTIo in our paper. After downloading the datasets, move them in
/path/to/flot/data/flownet3d/
Instead of moving the datasets, you can also create a link between the actual dataset location and the folder /path/to/flot/data/flownet3d/
:
$ ln -s /path/to/dataset/on/your/sytem /path/to/flot/data/flownet3d/
The directory /path/to/flot/data/flownet3d/
should have the following sub-directories:
/path/to/flot/data/flownet3d/data_processed_maxcut_35_20k_2k_8192 % FlyingThings3D dataset
/path/to/flot/data/flownet3d/kitti_rm_ground % KITTI dataset
- HPLFlowNet's version: Please follow the instructions given here to construct the FlyingThings3D and KITTI datasets used in this paper. These datasets are denoted FT3Ds and KITTIs in our paper. After preparation, move the datasets in
/path/to/flot/data/HPLFlowNet/
Instead of moving the datasets, you can also create a link between the actual dataset location and the folder /path/to/flot/data/HPLFlowNet/
:
$ ln -s /path/to/dataset/on/your/sytem /path/to/flot/data/HPLFlowNet/
The directory /path/to/flot/data/HPLFlowNet/
should have the following sub-directories:
/path/to/flot/data/HPLFlowNet/FlyingThings3D_subset_processed_35m % FlyingThings3D dataset
/path/to/flot/data/HPLFlowNet/KITTI_processed_occ_final % KITTI dataset
- Quickest test. Type:
$ cd /path/to/flot/scripts/
$ python val_test.py
This evaluates the model trained on FT3Do and 2048 points available at /path/to/flot/pretrained_models/model_2048.tar
on KITTIo.
For help on how to use this script, type:
$ cd /path/to/flot/scripts/
$ python val_test.py --help
- FlowNet3d's datasets. A model trained on FT3Do and 2048 points (Sec. 4.5 in our paper) is available at
/path/to/flot/pretrained_models/model_2048.tar
.
To evaluate this pre-trained model on KITTIo, type:
$ cd /path/to/flot/scripts/
$ python val_test.py --dataset flownet3D_kitti --test --nb_points 2048 --path2ckpt ../pretrained_models/model_2048.tar
To evaluate this pre-trained model on the test set of FT3Do, type:
$ cd /path/to/flot/scripts/
$ python val_test.py --dataset flownet3D_ft3d --test --nb_points 2048 --path2ckpt ../pretrained_models/model_2048.tar
- HPLFlowNet's datasets. A model trained on FT3Ds and 8192 points (Sec. 4.5 in our paper) is available at
/path/to/flot/pretrained_models/model_8192.tar
.
To evaluate this pretrained model on KITTIs, type:
$ cd /path/to/flot/scripts/
$ python val_test.py --dataset HPLFlowNet_kitti --test --nb_points 8192 --path2ckpt ../pretrained_models/model_8192.tar
To evaluate this pre-trained model on the test set of FT3Ds, type:
$ cd /path/to/flot/scripts/
$ python val_test.py --dataset HPLFlowNet_ft3d --test --nb_points 8192 --path2ckpt ../pretrained_models/model_8192.tar
By default, the model and tensorboard's logs are stored in /path/to/flot/experiments
. A script to train a flot model is available in /path/to/flot/train.py
. For help on how to use this script, please type:
$ cd /path/to/flot/scripts/
$ python train.py --help
- FlowNet3d's datasets. To train FLOT on the FlowNet3D's version of FlyingThing3D on 2048 points as in Sec. 4.5 of our paper, type:
$ cd /path/to/flot/
$ python train.py --nb_iter 1 --dataset flownet3d --nb_points 2048 --batch_size 4 --nb_epochs 400
These 400 epochs on 2048 points takes about 4 days to complete on a GeForce RTX 2080 Ti.
- HPLFlowNet's datasets. To train FLOT on the HPLFlownet's version of FlyingThing3D on 8192 points as in Sec. 4.4 of our paper, type:
$ cd /path/to/flot/
$ python train.py --nb_iter 1 --dataset HPLFlowNet --nb_points 8192 --batch_size 1 --nb_epochs 60
These 40 epochs on 8192 points takes about 6 days to complete on a GeForce RTX 2080 Ti.
It is possible to train FLOT on you own dataset by creating a new dataloader that inherits from flot.datasets.generic.SceneFlowDataset
.
Your new dataloader's class then needs to implement the function load_sequence(self, idx)
that loads the idx
example of the dataset. Please refer to the documentation of flot.datasets.generic.SceneFlowDataset.load_sequence
for more information.
Examples of dataloaders are available in the directory datasets
, see, e.g., flot.datasets.flyingthing3D_hplflownet
.
Once your new dataloader is implemented, it can be used in the script train.py
for training or in the val_test.py
for evaluation by importing this new dataset in the function my_main
.
Import FLOT by typing
from flot.models import FLOT
FLOT's constructor accepts one argument: nb_iter
, which is the number of unrolled iterations of the Sinkhorn algorithm. In our experiments, we tested 1, 3, and 5 iterations. For example:
flot = FLOT(nb_iter=3)
The simpler model FLOT0 is obtained by setting nb_iter=0
. In this case, the OT module is equivalent to an attention layer.
flot_0 = FLOT(nb_iter=0)
Input point clouds pc1
and pc2
can be passed to flot
to estimate the flow from pc1
to pc2
as follows:
scene_flow = flot([pc1, pc2])
The input point clouds pc1
and pc2
must be torch tensors of size batch_size x nb_points x 3
.
-
Currently a nearest neighbour search, needed to perform convolutions on points, is done by an exhaustive search in the function
flot.models.FLOT.forward
. This search could be moved in the dataloader where fast algorithm can be used. -
The transport cost
C
in the functionflot.tools.ot.sinkhorn
is computed densely, whereas several values are masked right after bysupport
. The computation ofsupport
could be done in the dataloader using fast nearest neighbours search. The computation ofC
could then be limited to the required entries. -
It is also possible to change the type convolution on point clouds used in
flot.models.FLOT.__init(self, nb_iter)__
. We provided an home-made, self-contained, implementation of PointNet++-like convolutions. One can use other convolutions on points for which faster implementation exists.
FLOT is released under the Apache 2.0 license.