Awesome
Learning Gradient Fields for Shape Generation
This repository contains a PyTorch implementation of the paper:
Learning Gradient Fields for Shape Generation [Project page] [Arxiv] [Short-video] [Long-video]
Ruojin Cai*, Guandao Yang*, Hadar Averbuch-Elor, Zekun Hao, Serge Belongie, Noah Snavely, Bharath Hariharan (* Equal contribution)
ECCV 2020 (Spotlight)
<p float="left"> <img src="assets/ShapeGF.gif" height="256"/> </p>Introduction
In this work, we propose a novel technique to generate shapes from point cloud data. A point cloud can be viewed as samples from a distribution of 3D points whose density is concentrated near the surface of the shape. Point cloud generation thus amounts to moving randomly sampled points to high-density areas. We generate point clouds by performing stochastic gradient ascent on an unnormalized probability density, thereby moving sampled points toward the high-likelihood regions. Our model directly predicts the gradient of the log density field and can be trained with a simple objective adapted from score-based generative models. We show that our method can reach state-of-the-art performance for point cloud auto-encoding and generation, while also allowing for extraction of a high-quality implicit surface.
Dependencies
# Create conda environment with torch 1.2.0 and CUDA 10.0
conda env create -f environment.yml
conda activate ShapeGF
# Compile the evaluation metrics
cd evaluation/pytorch_structural_losses/
make clean
make all
Dataset
Please follow the instruction from PointFlow to set-up the dataset: link.
Pretrained Model
Pretrained model will be available in the following google drive: link.
To use the pretrained models, download the pretrained
folder and put it under the project root directory.
Testing the pretrained auto-encoding model:
The following commands test the performance of the pre-trained models in the point cloud auto-encoding task. The commands output the CD and EMD on the test/validation sets.
# Usage:
# python test.py <config> --pretrained <checkpoint_filename>
python test.py configs/recon/airplane/airplane_recon_add.yaml \
--pretrained pretrained/recon/airplane_recon_add.pt
python test.py configs/recon/car/car_recon_add.yaml \
--pretrained pretrained/recon/car_recon_add.pt
python test.py configs/recon/chair/chair_recon_add.yaml \
--pretrained pretrained/recon/chair_recon_add.pt
The pretrained model's auto-encoding performance is as follows:
Dataset | Metrics | Ours | Oracle |
---|---|---|---|
Airplane | CD x1e4 | 0.966 | 0.837 |
EMD x1e2 | 2.632 | 2.062 | |
Chair | CD x1e4 | 5.660 | 3.201 |
EMD x1e2 | 4.976 | 3.297 | |
Car | CD x1e4 | 5.306 | 3.904 |
EMD x1e2 | 4.380 | 3.251 |
Testing the pretrained generation model:
The following commands test the performance of the pre-trained models in the point cloud generation task. The commands output the JSD, MMD-(CD/EMD), COV-(CD/EMD), and 1NN-(CD/EMD).
# Usage:
# python test.py <config> --pretrained <checkpoint_filename>
python test.py configs/gen/airplane_gen_add.yaml \
--pretrained pretrained/gen/airplane_gen_add.pt
python test.py configs/gen/car/car_gen_add.yaml \
--pretrained pretrained/gen/car_gen_add.pt
python test.py configs/gen/chair/chair_gen_add.yaml \
--pretrained pretrained/gen/chair_gen_add.pt
Training
Single GPU Training
# Usage:
python train.py <config>
Multi GPU Training
Our code also provides single-node multi GPU training using pytorch's Distributed Data Parallel. The script will run on all GPUs visible to the function. The usage and examples are as follows:
# Usage
python train_multi_gpus.py <config>
# To specify the total batch size, use --batch_size
python train_multi_gpus.py <config> --batch_size <#gpu x batch_size/GPU>
Stage-1: Auto-encoding
In this stage, we create a conditional generator that models the distribution of 3D points conditioned on the latent vector. The commands used to train our auto-encoding model for a single-shape, single ShapeNet category, and the whole ShapeNet are:
# Single shape
python train.py configs/recon/single_shapes/dress.yaml # the dress in the teaser
python train.py configs/recon/single_shapes/torus.yaml # the torus in the teaser
# Single category
python train.py configs/recon/airplane/airplane_recon_add.yaml # airplane
python train.py configs/recon/airplane/chair_recon_add.yaml # chair
python train.py configs/recon/airplane/car_recon_add.yaml # car
# Whole shape-net
python train_multi_gpus.py configs/recon/shapenet/shapenet_recon.yaml # ShapeNet
Stage-2: Generation
In the second stage, we train a l-GAN to model the distribution of shapes - which are captured by the latent vector of the auto-encoder described in the first stage.
The commands used to train l-GAN for a single ShapeNet category using the default pretrained model (in the <root>/pretrained
directory) are:
python train.py configs/gen/airplane_gen_add.yaml # airplane
python train.py configs/gen/chair_gen_add.yaml # chair
python train.py configs/gen/car_gen_add.yaml # car
Cite
Please cite our work if you find it useful:
@inproceedings{ShapeGF,
title={Learning Gradient Fields for Shape Generation},
author={Cai, Ruojin and Yang, Guandao and Averbuch-Elor, Hadar and Hao, Zekun and Belongie, Serge and Snavely, Noah and Hariharan, Bharath},
booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
year={2020}
}
Acknowledgment
This work was supported in part by grants from Magic Leap and Facebook AI, and the Zuckerman STEM leadership program.