📌 Description

SPT is a superpoint-based transformer 🤖 architecture that efficiently ⚡ performs semantic segmentation on large-scale 3D scenes. This method includes a fast algorithm that partitions 🧩 point clouds into a hierarchical superpoint structure, as well as a self-attention mechanism to exploit the relationships between superpoints at multiple scales.

📰 Updates

• 14.07.2023 Our paper Efficient 3D Semantic Segmentation with Superpoint Transformer was accepted at ICCV 2023 🥳
• 15.06.2023 Official release 🌱

💻 Environment requirements

This project was tested with:

• Linux OS
• NVIDIA GTX 1080 Ti 11G, NVIDIA V100 32G, NVIDIA A40 48G
• CUDA 11.8 (torch-geometric does not support CUDA 12.0 yet)
• conda 23.3.1

🏗 Installation

Simply run install.sh to install all dependencies in a new conda environment named spt.

# Creates a conda env named 'spt' env and installs dependencies
./install.sh

Note: See the Datasets page for setting up your dataset path and file structure.

🔩 Project structure

└── superpoint_transformer
    │
    ├── configs                   # Hydra configs
    │   ├── callbacks                 # Callbacks configs
    │   ├── data                      # Data configs
    │   ├── debug                     # Debugging configs
    │   ├── experiment                # Experiment configs
    │   ├── extras                    # Extra utilities configs
    │   ├── hparams_search            # Hyperparameter search configs
    │   ├── hydra                     # Hydra configs
    │   ├── local                     # Local configs
    │   ├── logger                    # Logger configs
    │   ├── model                     # Model configs
    │   ├── paths                     # Project paths configs
    │   ├── trainer                   # Trainer configs
    │   │
    │   ├── eval.yaml                 # Main config for evaluation
    │   └── train.yaml                # Main config for training
    │
    ├── data                      # Project data (see docs/datasets.md)
    │
    ├── docs                      # Documentation
    │
    ├── logs                      # Logs generated by hydra and lightning loggers
    │
    ├── media                     # Media illustrating the project
    │
    ├── notebooks                 # Jupyter notebooks
    │
    ├── scripts                   # Shell scripts
    │
    ├── src                       # Source code
    │   ├── data                      # Data structure for hierarchical partitions
    │   ├── datamodules               # Lightning DataModules
    │   ├── datasets                  # Datasets
    │   ├── dependencies              # Compiled dependencies
    │   ├── loader                    # DataLoader
    │   ├── loss                      # Loss
    │   ├── metrics                   # Metrics
    │   ├── models                    # Model architecture
    │   ├── nn                        # Model building blocks
    │   ├── optim                     # Optimization 
    │   ├── transforms                # Functions for transforms, pre-transforms, etc
    │   ├── utils                     # Utilities
    │   ├── visualization             # Interactive visualization tool
    │   │
    │   ├── eval.py                   # Run evaluation
    │   └── train.py                  # Run training
    │
    ├── tests                     # Tests of any kind
    │
    ├── .env.example              # Example of file for storing private environment variables
    ├── .gitignore                # List of files ignored by git
    ├── .pre-commit-config.yaml   # Configuration of pre-commit hooks for code formatting
    ├── install.sh                # Installation script
    ├── LICENSE                   # Project license
    └── README.md

Note: See the Datasets page for further details on `data/``.

Note: See the Logs page for further details on `logs/``.

🚀 Usage

Datasets

See the Datasets page to set up your datasets.

Evaluating SPT

Use the following commands to evaluate SPT from a checkpoint file checkpoint.ckpt:

# Evaluate SPT on S3DIS Fold 5
python src/eval.py experiment=s3dis datamodule.fold=5 ckpt_path=/path/to/your/checkpoint.ckpt

# Evaluate SPT on KITTI-360 Val
python src/eval.py experiment=kitti360  ckpt_path=/path/to/your/checkpoint.ckpt 

# Evaluate SPT on DALES
python src/eval.py experiment=dales ckpt_path=/path/to/your/checkpoint.ckpt

Note: The pretrained weights of the SPT and SPT-nano models for S3DIS 6-Fold, KITTI-360 Val, and DALES are available at:

Training SPT

Use the following commands to train SPT on a 32G-GPU:

# Train SPT on S3DIS Fold 5
python src/train.py experiment=s3dis datamodule.fold=5

# Train SPT on KITTI-360 Val
python src/train.py experiment=kitti360 

# Train SPT on DALES
python src/train.py experiment=dales

Use the following to train SPT on a 11G-GPU 💾 (training time and performance may vary):

# Train SPT on S3DIS Fold 5
python src/train.py experiment=s3dis_11g datamodule.fold=5

# Train SPT on KITTI-360 Val
python src/train.py experiment=kitti360_11g 

# Train SPT on DALES
python src/train.py experiment=dales_11g

Note: Encountering CUDA Out-Of-Memory errors 💀💾 ? See our dedicated troubleshooting section.

Note: Other ready-to-use configs are provided in configs/experiment/. You can easily design your own experiments by composing configs:

# Train Nano-3 for 50 epochs on DALES
python src/train.py datamodule=dales model=nano-3 trainer.max_epochs=50

See Lightning-Hydra for more information on how the config system works and all the awesome perks of the Lightning+Hydra combo.

Note: By default, your logs will automatically be uploaded to Weights and Biases, from where you can track and compare your experiments. Other loggers are available in configs/logger/. See Lightning-Hydra for more information on the logging options.

Notebooks & visualization

We provide notebooks to help you get started with manipulating our core data structures, configs loading, dataset and model instantiation, inference on each dataset, and visualization.

In particular, we created an interactive visualization tool ✨ which can be used to produce shareable HTMLs. Demos of how to use this tool are provided in the notebooks. Additionally, examples of such HTML files are provided in media/visualizations.7z

📚 Documentation

• README - General introduction to the project
• Data - Introduction to NAG and Data, the core data structures of this project
• Datasets - Introduction to Datasets and the project's data/ structure
• Logging - Introduction to logging and the project's logs/ structure

Note: We endeavoured to comment our code as much as possible to make this project usable. Still, if you find some parts are unclear or some more documentation would be needed, feel free to let us know by creating an issue !

👩‍🔧 Troubleshooting

Here are some common issues and tips for tackling them.

SPT on an 11G-GPU

Our default configurations are designed for a 32G-GPU. Yet, SPT can run on an 11G-GPU 💾, with minor time and performance variations.

We provide configs in configs/experiment/ for training SPT on an 11G-GPU 💾:

# Train SPT on S3DIS Fold 5
python src/train.py experiment=s3dis_11g datamodule.fold=5

# Train SPT on KITTI-360 Val
python src/train.py experiment=kitti360_11g 

# Train SPT on DALES
python src/train.py experiment=dales_11g

CUDA Out-Of-Memory Errors

Having some CUDA OOM errors 💀💾 ? Here are some parameters you can play with to mitigate GPU memory use, based on when the error occurs.

💳 Credits

• This project was built using Lightning-Hydra template.
• The main data structures of this work rely on PyToch Geometric
• Some point cloud operations were inspired from the Torch-Points3D framework, although not merged with the official project at this point.
• For the KITTI-360 dataset, some code from the official KITTI-360 was used.
• Some superpoint-graph-related operations were inspired from Superpoint Graph
• The hierarchical superpoint partition is computed using Parallel Cut-Pursuit

Citing our work

If your work uses all or part of the present code, please include the following a citation:

@inproceedings{robert2023spt,
  title={Efficient 3D Semantic Segmentation with Superpoint Transformer},
  author={Robert, Damien and Raguet, Hugo and Landrieu, Loic},
  journal={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  year={2023}
}

You can find our paper on arxiv 📄.