davidgillsjo / bssc-net

This is a PyTorch implementation of a Bayesian Convolutional Neural Network (BCNN) for Semantic Scene Completion on the SUNCG dataset. Given a depth image the network outputs a semantic segmentation and entropy score in 3D voxel format.

BSSC-Net

CNN for Bayesian Semantic Scene Completion from a single depth image. This is the code release for our paper from ICPR2020. Preprint is also available at ArXiV.

Introduction

This is a PyTorch implementation of a Bayesian Convolutional Neural Network (BCNN) for Semantic Scene Completion on the SUNCG dataset. Given a depth image represented as inverted depth (see code) the network outputs a semantic segmentation and entropy score in 3D voxel format.

Setup repository

git clone [email protected]:DavidGillsjo/bssc-net.git --recurse-submodules -j8

Using the docker images

Check your Nvidia GPU model and adjust the variable gpu_arch accordingly in docker/<module>/Dockerfile.

With elevated docker group permissions:

cd docker/<module>
./build.sh
./run.sh

Without:

cd docker/<module>
sudo DUSER=<username> ./build.sh
sudo DHOME=/home/<username> ./run.sh

Render and Build Dataset

Download SUNCG dataset

Download camera viewpoints:

wget http://pbrs.cs.princeton.edu/pbrs_release/data/camera_v2.zip
unzip camera_v2.zip -d ../camera

Build and run the suncg_house3d docker in the docker folder, see section Using the docker images. Alternatively install the dependencies yourself, see the Dockerfile.

Go to repo root from docker:

cd /host_home/<your-repo-path>

the following code snippets assume your are in this directory.

Build SUNCG Toolbox:

cd libs/SUNCGtoolbox/gaps
make clean
make

Build House3d:

cd libs/House3D/renderer
PYTHON_CONFIG=python3-config make -j

Set python path:

. init_env.sh

Execute rendering and grid generation script:

cd preprocessing
python3 generate_grids.py <suncg_dir> --nbr-proc 8 --model-blacklist blacklists/default.yaml --datasets <data_split_dir>/*mini.json

This will store results in <suncg_dir>/scene_comp, for more details see python3 generate_grids.py --help. JSON dataset files are generated with preprocessing/generate_data_splits.py. The splits used in the article can be found here.

Prerequisites for training/evaluating networks

Build and run the bssc docker in the docker folder, see section Using the docker images. Alternatively install the dependencies yourself, see the Dockerfile.

Train network

Depth images with corresponding voxel ground truth generated in previous section is required. For example placed in /host_home/<suncg_dir>/scene_comp.

Training can be monitored with tensorboard, for this you need to run:

cd ssc/scripts
./start_tensorboard.sh &

To start training, run:

cd ssc/scripts
python3 train.py /host_home/<suncg_dir>/scene_comp <train_json> --cfg ../cfg/train_bayesian.yaml --val <val_json>

See configs for options.

Evaluate example image

To run the example, first download the pre-trained weights from here. Then run the evaluation script:

cd ssc/scripts
python3 eval.py ../../example ../../example/dataset.json bssc.tar --cfg ../cfg/eval_bayesian.yaml --result-dir <my_result_path>

For the SSC style net or run the following for Unet implementation:

cd ssc/scripts
python3 eval.py ../../example ../../example/dataset.json bssc_unet.tar --cfg ../cfg/eval_bayesian_unet.yaml --result-dir <my_result_path>

Citation

@INPROCEEDINGS{bssc_net,
  author={Gillsjö, David and Åström, Kalle},
  booktitle={2020 25th International Conference on Pattern Recognition (ICPR)}, 
  title={In Depth Bayesian Semantic Scene Completion}, 
  year={2021},
  volume={},
  number={},
  pages={6335-6342},
  doi={10.1109/ICPR48806.2021.9412403}}

Acknowledgment

This work is supported by Wallenberg AI Autonomous Systems and Software Program.