PointNet++ Semantic Segmentation on ScanNet in PyTorch with CUDA acceleration based on the original PointNet++ repo and the PyTorch implementation with CUDA
The semantic segmentation results in percentage on the ScanNet train/val split in data/.
| use XYZ | use color | use normal | use multiview | use MSG | mIoU | weights |
|---|---|---|---|---|---|---|
| ✔️ | ✔️ | ✔️ | - | - | 50.48 | download |
| ✔️ | ✔️ | ✔️ | - | ✔️ | 52.50 | download |
| ✔️ | - | ✔️ | ✔️ | - | 65.75 | download |
| ✔️ | - | ✔️ | ✔️ | ✔️ | 67.60 | download |
If you want to play around with the pre-trained model, please download the zip file and unzip it under outputs/.
- Linux (tested on Ubuntu 14.04/16.04)
- Python 3.6+
- PyTorch 1.8
- TensorBoardX
Please run conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch to install PyTorch 1.8 and run pip install -r requirements.txt to install other required packages.
Install this library by running the following command:
cd pointnet2
python setup.py installChange the path configurations for the ScanNet data in lib/config.py
-
Download the ScanNet frames here (~13GB) and unzip it under the project directory.
-
Extract the multiview features from ENet:
python scripts/compute_multiview_features.py- Generate the projection mapping between image and point cloud
python scripts/compute_multiview_projection.py- Project the multiview features from image space to point cloud
python scripts/project_multiview_features.pyNote you might need ~100GB RAM to train the model with multiview features
Parse the ScanNet data into *.npy files and save them in preprocessing/scannet_scenes/
python preprocessing/collect_scannet_scenes.pyDon't forget to visualize the preprocessed scenes to check the consistency
python preprocessing/visualize_prep_scene.py --scene_id <scene_id>The visualized <scene_id>.ply is stored in preprocessing/label_point_clouds/ - Drag that file into MeshLab and you'll see something like this:
Train the PointNet++ semantic segmentation model on ScanNet scenes with raw RGB values and point normals (for more training options, see python scripts/train.py -h)
python scripts/train.py --use_color --use_normal --use_msgThe trained models and logs will be saved in outputs/<time_stamp>/
Evaluate the trained models and report the segmentation performance in point accuracy, voxel accuracy and calibrated voxel accuracy
python scripts/eval.py --folder <time_stamp>Note that all model options must match the ones used for training.
Visualize the semantic segmentation results on points in a given scene
python scripts/visualize.py --folder <time_stamp> --scene_id <scene_id>Note that all model options must match the ones used for training.
The generated <scene_id>.ply is stored in outputs/<time_stamp>/preds - Drag that file into MeshLab and you'll see something like the one below. See the class palette here
- 07/29/2021 Upgrade to PyTorch 1.8 & fix existing issues
- 03/29/2020 Release the code
- Release all pretrained models
- Upgrade to PyTorch 1.8
- Fix issues with loading pre-trained models
- charlesq34/pointnet2: Paper author and official code repo.
- erikwijmans/Pointnet2_PyTorch: Initial work of PyTorch implementation of PointNet++ with CUDA acceleration.