Skip to content
Switch branches/tags

3D Point Capsule Networks

Created by Yongheng Zhao, Tolga Birdal, Haowen Deng, Federico Tombari from TUM.

This repository contains the implementation of our CVPR 2019 paper 3D Point Capsule Networks. In particular, we release code for training and testing a 3D-PointCapsNet network for classification, reconstruction, part interpolation and extraction of 3d local descriptors as well as the pre-trained models for quickly replicating our results.

For an intuitive explanation of the 3D point capsule networks, please check out Tolga's CVPR tutorial.

Part Interpolation

Distribution of capsule reconstruction


In this paper, we propose 3D point capsule networks, an auto-encoder designed to process sparse 3D point clouds while preserving spatial arrangements of the input data. 3D capsule networks arise as a direct consequence of our novel unified 3D auto-encoder formulation. Their dynamic routing scheme and the peculiar 2D latent space deployed by our approach bring in improvements for several common point cloud-related tasks, such as object classification, object reconstruction and part segmentation as substantiated by our extensive evaluations. Moreover, it enables new applications such as part interpolation and replacement.


If you find our work useful in your research, please consider citing:

		author={Zhao, Yongheng and Birdal, Tolga and Deng, Haowen and Tombari, Federico}, 
		booktitle={Conference on Computer Vision and Pattern Recognition (CVPR)}, 
		title={3D Point 
		Capsule Networks}, 


The code is based on PyTorch. It has been tested with Python 3.6, PyTorch 1.0.0, CUDA 9.2(or higher) on Ubuntu 16.04. (You can also use PyTorch 0.4.0 but you need to replace the Chamfer distance package with the original nndistance.)

Install h5py for Python:

  sudo apt-get install libhdf5-dev
  sudo pip install h5py

Install Chamfer Distance(CD) package: (Be aware of the PyTorch 1.0.1. It may have a problem for building this cuda package.)

  cd models/nndistance
  python install

(In case you are using pytorch version higher than 1.0, you could use the updated chamfer distance package named "torch-nndistance". But you need to modify the package usage in the several scripts in which the CD library is used. You can find "" in the updated package folder for the usage reference.

To visualize the training process in PyTorch, consider installing TensorBoard.

To visualize the reconstructed point cloud, consider installing Open3D.


ShapeNetPart Dataset

  cd dataset

ShapeNet Core with 13 categories (refered from AtlasNet.)

  cd dataset

ShapeNet Core with 55 categories (refered from FoldingNet.)

  cd dataset

Pre-trained model

You can download the pre-trained models here.


A Minimal Example

We provide an example demonstrating the basic usage in the folder 'mini_example'.

To visualize the reconstruction from latent capsules with our pre-trained model:

  cd mini_example/AE
  python --model ../../checkpoints/shapenet_part_dataset_ae_200.pth

To train a point capsule auto encoder with ShapeNetPart dataset by yourself:

  cd mini_example/AE

Point Capsule Auto Encoder

To train a point capsule auto encoder with another dataset:

  cd apps/AE
  python --dataset < shapenet_part, shapenet_core13, shapenet_core55 >

To monitor the training process, use TensorBoard by specifying the log directory:

  tensorboard --logdir log

To test the reconstruction accuracy:

  python  --dataset < >  --model < >
  python --dataset shapenet_core13 --model ../../checkpoints/shapenet_core13_dataset_ae_230.pth

To visualize the reconstructed points:

  python --dataset < >  --model < >
  python --dataset shapenet_core13 --model ../../checkpoints/shapenet_core13_dataset_ae_230.pth

Transfer Learning and Semi Supervised Classification

To generate latent capsules from a pre-trained model and save them into a file:

  cd apps/trasfer_cls
  python --dataset < >  --model < >  --save_training  # process and save training dataset
  python --dataset < >  --model < >   # process and save testing dataset

To train and test the liner SVM classifier with the pre-trained AE model and pre-saved latent capsules:

  python --dataset < >  --model < >

The AE model and latent capsules are obtained from different datasets in order to test the performance of classification under transfer.

Training a Liner SVM classifier with a limited part of the training data and testing with the complete test data:

  python --dataset < >  --model < > --percent_training_dataset < 5, 10 ...>
  python --dataset shapenet_part  --model ../../checkpoints/shapenet_part_dataset_ae_200.pth --percent_training_dataset 10

Part Segmentation

To generate latent capsules with the part label from a pre-trained model and save them into a file (The model is also trained with shapenet-part dataset):

  cd apps/part_seg
  python --dataset shapenet_part  --model < >  --save_training  # process and save training dataset
  python --dataset shapenet_part  --model < >   # process and save testing dataset

To train a capsule-wise part segmentation with a specific amount of training data:

  python --model < > --percent_training_dataset < 5, 10 ...>
  python --model ../../checkpoints/shapenet_part_dataset_ae_200.pth --percent_training_dataset 1

To evaluate and visualize the part segmentation: (--model < pre-trained model of point capsule auto encoder >; --part_model )

 python --model < >  --part_model < > --class_choice < >
 python --model ../../checkpoints/shapenet_part_dataset_ae_200.pth  --part_model ../../checkpoints/part_seg_1percent.pth  --class_choice Airplane

Part Interpolation and Replacement

To visualize the part interpolation in open3D:

 python --model < >  --part_model < > --class_choice < >
 python --model ../../checkpoints/shapenet_part_dataset_ae_200.pth  --part_model ../../checkpoints/part_seg_100percent.pth  --class_choice Airplane

To visualize the part replacement in open3D:

 python --model < >  --part_model < > --class_choice < >

3D Local Feature Extraction

to be continued...


Our code is released under MIT License (see LICENSE file for details).


The chamfer distance package is based on nndistance. The necessary modifications have been done to this repository in order to run it with PyTorch 1.0.0.

The capsule layer is based upon and modified from Capsule-Network-Tutorial

Our capsule decoder is based upon the decoder of AtlasNet.