ShapeInversion

Paper

Junzhe Zhang, Xinyi Chen, Zhongang Cai, Liang Pan, Haiyu Zhao, Shuai Yi, Chai Kiat Yeo, Bo Dai, Chen Change Loy "Unsupervised 3D Shape Completion through GAN Inversion" CVPR 2021

Project page: link
Paper: link

Results

Setup

Environment

conda create -n shapeinversion python=3.7
conda activate shapeinversion
pip install torch==1.2.0 torchvision==0.4.0
pip install plyfile h5py Ninja matplotlib scipy

Datasets

Our work is extensively evaluated with several existing datasets. For the virtual scan benchmark (derived from ShapeNet), we use CRN's dataset. We would suggest you to get started with this dataset. For ball-holed partial shapes, we refer to PF-Net. For PartNet dataset, we download from MPC. For real scans processed from KITTI, MatterPort3D, and ScanNet, we get from pcl2pcl.

Get started

We provide pretrained tree-GAN models for you to directly start with the inversion stage. You can download them from Google drive or Baidu cloud (password: w1n9), and put them to the pretrained_models folder.

Shape completion

You can specify other class and other datasets, like real scans provided by pcl2pcl.

python trainer.py \
--dataset CRN \
--class_choice chair \
--inversion_mode completion \
--mask_type k_mask \
--save_inversion_path ./saved_results/CRN_chair \
--ckpt_load pretrained_models/chair.pt \
--dataset_path <your_dataset_directory>

Evaluating completion results

For datasets with GT, such as the above CRN_chair:

python eval_completion.py \
--eval_with_GT true \
--saved_results_path saved_results/CRN_chair

For datasets without GT:

python eval_completion.py \
--eval_with_GT false \
--saved_results_path <your_results_on_KITTI>

Giving multiple valid outputs

ShapeInversion is able to provide multiple valid complete shapes, especially when extreme incompleteness that causes ambiguity.

python trainer.py \
--dataset CRN \
--class_choice chair \
--inversion_mode diversity \
--save_inversion_path ./saved_results/CRN_chair_diversity \
--ckpt_load pretrained_models/chair.pt \
--dataset_path <your_dataset_directory>

Shape jittering

ShapeInversion is able to change an object into other plausible shapes of different geometries.

python trainer.py \
--dataset CRN \
--class_choice plane \
--save_inversion_path ./saved_results/CRN_plane_jittering  \
--ckpt_load pretrained_models/plane.pt \
--inversion_mode jittering \
--iterations 30 30 30 30 \
--dataset_path <your_dataset_directory>

Shape morphing

ShapeInversion enables morphing between two shapes.

python trainer.py \
--dataset CRN \
--class_choice chair \
--save_inversion_path ./saved_results/CRN_chair_morphing  \
--ckpt_load pretrained_models/chair.pt \
--inversion_mode morphing \
--dataset_path <your_dataset_directory>

Pretraining

You can also pretrain tree-GAN by yourself.

python pretrain_treegan.py \
--split train \
--class_choice chair \
--FPD_path ./evaluation/pre_statistics_chair.npz \
--ckpt_path ./pretrain_checkpoints/chair \
--knn_loss True \
--dataset_path <your_dataset_directory>

NOTE:

The inversion stage supports distributed training by simply adding --dist. It is tested on slurm as well.
The hyperparameters provided may not be optimal, feel free to tune them.
Smaller batch size for pretraining is totally fine.

Acknowledgement

The code is in part built on tree-GAN and DGP. Besides, CD and EMD are borrowed from ChamferDistancePytorch and MSN respectively, both of which are included in the external folder for convenience.

Citation

@inproceedings{zhang2021unsupervised,
    title = {Unsupervised 3D Shape Completion through GAN Inversion},
    author = {Zhang, Junzhe and Chen, Xinyi and Cai, Zhongang and Pan, Liang and Zhao, Haiyu 
    and Yi, Shuai and Yeo, Chai Kiat and Dai, Bo and Loy, Chen Change},
    booktitle = {CVPR},
    year = {2021}}

junzhezhang / shape-inversion