Skip to content

ashawkey/bsp_cvae

Repository files navigation

BSP-CVAE

This repo implements a BSP-CVAE model, which uses the idea of BSP-Net but is a generative model.

  • reconstruction on ScanNet:

recons

  • generation (interpolated latent codes):

generate

Install

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

Datasets

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.

Train

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.

Test

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.

References

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors