This repo implements a BSP-CVAE model, which uses the idea of BSP-Net but is a generative model.
- reconstruction on ScanNet:
- generation (interpolated latent codes):
This repo is tested on Ubuntu16.04, CUDA 10.1
.
For the python dependencies, see requirements.txt
.
We also use two Cython
extensions, install them by python setup.py build_ext --inplace
.
conda env create -n bsp
conda activate bsp
pip install -r requirements.txt
python setup.py build_ext --inplace
Data are assumed to be located in ${data_root}/datasets/
, where ${data_root}
can be set in main.py
.
-
ShapeNet
We use the preprocessed data provided by RfDNet, please follow their instructions and put it under${data_root}/datasets/ShapeNetv2_data
. We use 8 classes ('table', 'chair', 'bookshelf', 'sofa', 'trash_bin', 'cabinet', 'display', 'bathtub'
) in ShapeNet to train the model. -
ScanNet & Scan2CAD
If you want to test the reconstruction performance under indoor scenes, you should also download preprocessed ScanNet and Scan2CAD datasets following these instructions, and put it under${data_root}/datasets/scannet
.
The options and parameters should be modified directly in main.py
.
# train with default settings.
bash train.sh
It takes about 4 days to train the model for 800 epochs on a single GPU.
Logs are saved in workspace/log_${exp_name}.txt
.
Checkpoints are saved in workspace/checkpoints
.
Tensorboard records are saved in workspace/run
.
After training, just run the tests as follows:
# reconstruction on shapenet
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --test_shapenet
# generation
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --generate
# interpolated generation
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --interpolated_generate
# reconstruction on scannet
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --test_scannet
# save zs
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --save_z
# save database
CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python main.py --checkpoint latest --save_db
Results are saved in workspace/results
.