This repository includes implementations to the method presented in the paper "Shape-Pose Disentanglement using SE(3)-equivariant Vector Neurons", including additional support for losses types and architectures (see code).
We use the ShapeNet dataset as in AtlasNetV2. Please follow the instructions at CanonicalCapsules to convert the data to h5 files. Other datasets will be supported shortly.
The Dockerfile and the requirements.txt file include all the required dependencies. Please build and run it.
git clone VNT
cd VNT
docker build -f Dockerfile -t vnt --build-arg USER_ID=$(id -u) --build-arg GROUP_ID=$(id -g)
docker run -t -d --gpus all -v ${PWD}:/app/VNT --name vnt_experiment -p 8097:8097 vnt
Execute the running container:
docker exec -it vnt_experiment /bin/bash
cd VNT
To view training results and loss plots, run
python -m visdom.server
and click the URL http://localhost:8097/
Once visdom is up, start training
python train.py --model shape_pose --name chairs_paper --dataroot datasets/shapenet --class_choice chair --add_noise --apply_can_rot_loss --remove_knn 100 --resample --fps
An improved version of defaults can be seen in shape_pose2 model:
python train.py --model shape_pose2 --name chairs_updated --dataroot datasets/shapenet --class_choice chair --add_noise --apply_can_rot_loss --remove_knn 100 --resample --fps
python test.py --model shape_pose --name chairs_paper --dataroot datasets/shapenet --class_choice chair
or
python test.py --model shape_pose2 --name chairs_updated --dataroot datasets/shapenet --class_choice chair
Pre-trained model can be found here
Download the models and place the extracted folder under the checkpoints directory.
@article{katzir2022shape,
title={Shape-Pose Disentanglement using SE (3)-equivariant Vector Neurons},
author={Katzir, Oren and Lischinski, Dani and Cohen-Or, Daniel},
journal={arXiv preprint arXiv:2204.01159},
year={2022}
}
MIT License
The structure of this codebase is borrowed from the pytorch implementation of CycleGAN-pix2pix.
The code includes parts of the source code from VNN and AtlasNet2