PQ-NET

This repository provides PyTorch implementation of our paper:

PQ-NET: A Generative Part Seq2Seq Network for 3D Shapes

Rundi Wu, Yixin Zhuang, Kai Xu, Hao Zhang, Baoquan Chen

CVPR 2020

Prerequisites

• Linux
• NVIDIA GPU + CUDA CuDNN
• Python 3.6

Dependencies

Install python package dependencies through pip:

pip install -r requirements.txt

Compile the extension module brought from Occupancy_Networks:

python setup.py build_ext --inplace

Data

We first voxelized PartNet shapes and scale each part to $64^3$ resolution. We provide data for three categories: chair, table, lamp. Please use link1(PKU disk) or link2(Google Drive) to download the voxelized PartNet shapes and exact the file to data/ folder, e.g.

cd data
tar -xvf Lamp.tar.gz

Then run data/sample_points_from_voxel.py to sampled paired points and signed values, e.g:

python data/sample_points_from_voxel.py --src data --category Lamp

Training

Example training scripts can be found in scripts folder. See config/ for specific definition of all hyper-parameters.

To train the main model:

# train part auto-encoder following multiscale strategy 16^3-32^3-64^3
sh scripts/lamp/train_lamp_partae_multiscale.sh # use two gpus 

# train seq2seq model
sh scripts/lamp/train_lamp_seq2seq.sh

For random generation task, further train a latent GAN:

# encode each shape to latent space
sh scripts/lamp/enc_lamp_seq2seq.sh

# train latent GAN (wgan-gp)
sh scripts/lamp/train_lamp_lgan.sh

The trained models and experment logs will be saved in proj_log/pqnet-PartNet-Lamp/ by default.

Testing

Example testing scripts can also be found in scripts folder.

• Shape Auto-encoding

  After training the main model, run the model to reconstruct all test shapes:

  sh scripts/lamp/rec_lamp_seq2seq.sh

• Shape Generation

  After training the latent GAN, run latent GAN and the main model to do random generation:

  # run latent GAN to generate fake latent vectors
  sh scripts/lamp/test_lamp_lgan.sh
  
  # run the main model to decode the generated vectors to final shape
  sh scripts/lamp/dec_lamp_seq2seq.sh

The results will be saved inproj_log/pqnet-PartNet-Lamp/results/ by default.

Pre-trained models

Please use link1(PKU disk) or link2(Google Drive) to download the pretrained model. Download and extract it under proj_log/, so that all test scripts can be directly excecuted.

Voxelization

For those who need to train the model on their own dataset, see the instructions and code of our voxelization process here.

Cite

Please cite our work if you find it useful:

@InProceedings{Wu_2020_CVPR,
author = {Wu, Rundi and Zhuang, Yixin and Xu, Kai and Zhang, Hao and Chen, Baoquan},
title = {PQ-NET: A Generative Part Seq2Seq Network for 3D Shapes},
booktitle = {IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
month = {June},
year = {2020}
}