Skip to content


Repository files navigation


PyTorch implementation of our CVPR 2019 paper:

Single-Image Piece-wise Planar 3D Reconstruction via Associative Embedding

Zehao Yu*, Jia Zheng*, Dongze Lian, Zihan Zhou, Shenghua Gao

(* Equal Contribution)

Getting Started


Clone repository and use git-lfs to fetch the trained model (or download here):

git clone

We use Python 3. Create an Anaconda enviroment and install the dependencies:

conda create -y -n plane python=3.6
conda activate plane
conda install -c menpo opencv
pip install -r requirements.txt

Downloading and converting data

Please download the .tfrecords files for training and testing converted by PlaneNet, then convert the .tfrecords to .npz files:

python data_tools/ --data_type=train --input_tfrecords_file=/path/to/planes_scannet_train.tfrecords --output_dir=/path/to/save/processd/data
python data_tools/ --data_type=val --input_tfrecords_file=/path/to/planes_scannet_val.tfrecords --output_dir=/path/to/save/processd/data


Run the following command to train our network:

python train with dataset.root_dir=/path/to/save/processd/data


Run the following command to evaluate the performance:

python eval with dataset.root_dir=/path/to/save/processd/data resume_dir=/path/to/ dataset.batch_size=1


Run the following command to predict on a single image:

python with image_path=/path/to/image


We thank Chen Liu for his great works and repos.


Please cite our paper for any purpose of usage.

  author    = {Zehao Yu and Jia Zheng and Dongze Lian and Zihan Zhou and Shenghua Gao},
  title     = {Single-Image Piece-wise Planar 3D Reconstruction via Associative Embedding},
  booktitle = {CVPR},
  pages     = {1029--1037},
  year      = {2019}