Learning to Segment 3D Point Clouds in 2D Image Space

This repository is a community implementation of the paper Learning to Segment 3D Point Clouds in 2D Image Space by Lyu et al. using the PyTorch framework. The official TensorFlow implementation by the authors can be found here.

Getting started

Install environment

Install the conda environment with all necessary packages by

conda env create -f environment.yml

and activate it by

conda activate 3d-seg-in-2d

Download ShapeNet Part dataset

./download_data.sh

will download and extract the ShapeNet Part dataset from Stanford University into the directory hdf5_data.

Project 3D point clouds into 2D image space

python prepare_dataset.py

Training

Create a custom configuration file custom_config.yaml or use the default config.yaml to determine the hyperparameters and other settings. Run

python train_unet.py --config_path custom_config.yaml

or simply python train_unet.py to use the default configuration.

Configuration options

• device: Enable CUDA acceleration by setting the device in the configuration to device: 'cuda:0' or similar. MPS for Apple Silicon is also supported, set the device to device: 'mps'.

Evaluation

python test.py --model_path <MODEL_PATH>

to evaluate the model located at <MODEL_PATH>.

Result visualization

To visualize the segmentation result for class <idx_class> and sample <idx_class_sample>, run

python visualization.py <idx_class> <idx_class_sample>

Contributing

You can run the pre-commit hooks locally with pre-commit run.