PlankAssembly: Robust 3D Reconstruction from Three Orthographic Views with Learnt Shape Programs

Wentao Hu* · Jia Zheng* · Zixin Zhang* · Xiaojun Yuan · Jian Yin · Zihan Zhou

IEEE/CVF Conference on Computer Vision (ICCV), 2023

These authors contributed equally to this work.

Note

This branch contains the implementation of comparison method, PolyGen. Similar to our method, PolyGen adopts a Transformer-based architecture and proceeds by generating a set of 3D vertices, which are then connected to form 3D faces.

Setup

Our code has been tested with Python 3.8, PyTorch 1.10.0, CUDA 11.3 and PyTorch Lightning 1.7.6.

Please follow the instructions in the main branch to set up your environment and download the dataset.

Data Processing

To train PolyGen, we first parse our shape program into vertex sequence and face sequence:

python dataset/preprocess.py

Training

Use the following command to train vertex model and face model from scratch:

# train vertex model
python vertex_trainer.py fit --config configs/vertex.yaml
# train face model
python face_trainer.py fit --config configs/face.yaml

Testing

Use the following command to test with our pretrained model (vertex model weight, face model weight) or your own checkpoint:

# please first modify the checkpoint path in configs/test.yaml
python tester.py test --config configs/test.yaml

Visualization

To visualize the results, we build 3D mesh models from predictions:

python postprocess.py --exp_path path/to/lightning_log/dir

LICENSE

PlankAssembly is licensed under the AGPL-3.0 license. The code snippets in the third_party are available under Apache-2.0 License.