This is the official repository for the PyTorch implementation of "Canonical Capsules: Self-Supervised Capsules in Canonical Pose" by Weiwei Sun*, Andrea Tagliasacchi*, Boyang Deng, Sara Sabour, Soroosh Yazdani, Geoffrey Hinton, Kwang Moo Yi.

• Project Website
• PDF (arXiv)
• PDF (github copy)

⚠️ If you use this source core or data in your research (in any shape or format), we require you to cite our paper as:

@conference{sun2020canonical,
   title={Canonical Capsules: Self-Supervised Capsules in Canonical Pose},
   author={Weiwei Sun and Andrea Tagliasacchi and Boyang Deng and 
           Sara Sabour and Soroosh Yazdani and Geoffrey Hinton and
           Kwang Moo Yi},
   booktitle={Neural Information Processing Systems},
   year={2021}
}

Please install dependencies with the provided environment.yml:

conda env create -f environment.yml

• We use the ShapeNet dataset as in AtlasNetV2: download the data from AtlasNetV2's official repo and convert the downloaded data into h5 files with the provided script (i.e., data_utils/ShapeNetLoader.py).

• For faster experimentation, please use our 2D planes dataset, which we generated from ShapeNet (please cite both our paper, as well as ShapeNet if you use this dataset).

To train the model on 2D planes (training of network takes only 50 epochs, and one epoch takes approximately 2.5 minutes on an NVIDIA GTX 1080 Ti):

./main.py --log_dir=plane_dim2 --indim=2 --scheduler=5

To visualize the decompostion and reconstruction:

./main.py --save_dir=gifs_plane2d --indim=2 --scheduler=5 --mode=vis --pt_file=logs/plane_dim2/checkpoint.pth

To train the model on the 3D dataset:

./main.py --log_dir=plane_dim3 --indim=3 --cat_id=-1

We test the model with:

./main.py --log_dir=plane_dim3 --indim=3 --cat_id=-1 --mode=test

Note that the option cat_id indicates the category id to be used to load the corresponding h5 files (this look-up table):

We release the 3D pretrained models for both single categy (airplanes), as well as multi-category (all 13 classes).

To use our classification script:

python classification.py --data_dir=/path/to/saved/features --feature_type=caca --method_type=svm --use_kpts