This is the official code of

Fast Point Cloud Generation with Straight Flows
Lemeng Wu, Dilin Wang, Chengyue Gong, Xingchao Liu, Yunyang Xiong, Rakesh Ranjan, Raghuraman Krishnamoorthi, Vikas Chandra, Qiang Liu

Now we release code for training and inference. Some works are still in progress including pretrained checkpoint.

This code is largely build based on PVD. Make sure at least the following environments are installed (newer version may also works, we test in the below environments).

python==3.8
pytorch==1.4.0
torchvision==0.5.0
cudatoolkit==10.1
matplotlib==2.2.5
tqdm==4.32.1
open3d==0.9.0
trimesh=3.7.12
scipy==1.5.1

We also need to install pytorch3D for Chamfer Distance Loss, we recommend to follow the offical install guideline here

Install PyTorchEMD by

cd metrics/PyTorchEMD
python setup.py install
cp build/**/emd_cuda.cpython-36m-x86_64-linux-gnu.so .

We use the data follow PVD and PointFlow, which can be downloaded here. Extract and put the data in ./data/ folder/

First Stage, train the flow model. We do not add EMA here for a simple and quick converge as illustration.

$ python train_flow.py --category car|chair|airplane

Assume the checkpoint is saved as flow_checkpoint.pth (you can find it in the ./output/train_flow/ )

Second Stage, straight the flow, first sample the data pairs. We provide a single GPU version, in practice, we use multiGPU to speed up.

$ python sample_flow.py --category car|chair|airplane --model flow_checkpoint.pth

Then run the reflow procedure:

$ python train_reflow.py --category car|chair|airplane --model flow_checkpoint.pth

Assume the checkpoint is saved as reflow_checkpoint.pth (you can find it in the ./output/train_reflow/ )

Third Stage, distill the flow.

$ python train_distill.py --category car|chair|airplane --model reflow_checkpoint.pth

Assume the checkpoint is saved as distill_checkpoint.pth (you can find it in the ./output/train_distill/ )

$ python test_flow.py --category car|chair|airplane --model {flow|reflow|distill}_checkpoint.pth --step 1|20|50|100|500|1000

You can adjust the step in this test code. For flow, reflow model, we can still expect a good few-step generation.

@InProceedings{Wu_2023_CVPR,
    author    = {Wu, Lemeng and Wang, Dilin and Gong, Chengyue and Liu, Xingchao and Xiong, Yunyang and Ranjan, Rakesh and Krishnamoorthi, Raghuraman and Chandra, Vikas and Liu, Qiang},
    title     = {Fast Point Cloud Generation With Straight Flows},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {9445-9454}
}

This code is built based on PVD. Thanks for their great code repo!