PlanarReconstruction

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)

<img src="misc/pipeline.jpg" width="800">

Getting Started

Installation

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

git clone git@github.com:svip-lab/PlanarReconstruction.git

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/convert_tfrecords.py --data_type=train --input_tfrecords_file=/path/to/planes_scannet_train.tfrecords --output_dir=/path/to/save/processd/data
python data_tools/convert_tfrecords.py --data_type=val --input_tfrecords_file=/path/to/planes_scannet_val.tfrecords --output_dir=/path/to/save/processd/data

Training

Run the following command to train our network:

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

Evaluation

Run the following command to evaluate the performance:

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

Prediction

Run the following command to predict on a single image:

python predict.py with resume_dir=pretrained.pt image_path=/path/to/image

Acknowledgements

We thank Chen Liu for his great works and repos.

Citation

Please cite our paper for any purpose of usage.

@inproceedings{YuZLZG19,
  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}
}