ShapeFormer: Transformer-based Shape Completion via Sparse Representation

Project Page | Paper (ArXiv) | Twitter thread

This repository is the official pytorch implementation of our paper, ShapeFormer: Transformer-based Shape Completion via Sparse Representation.

Xinggaung Yan¹, Liqiang Lin¹, Niloy Mitra², Dani Lischinski³, Danny Cohen-Or⁴, Hui Huang^1†
1Shenzhen University, ²University College London, ³Hebrew University of Jerusalem, ⁴Tel Aviv University

⏳ UPDATES

• Core model code is released, please check core_code/README.md
• The complete code is released! Please have a try!
• (DFAUST) The data preprocess code for D-FAUST human shape is released!
• Add Google Colab

Installation

The code is tested in docker enviroment pytorch/pytorch:1.6.0-cuda10.1-cudnn7-devel. The following are instructions for setting up the environment in a Linux system from scratch. You can also directly pull our provided docker environment: sudo docker pull qheldiv/shapeformer Or build the docker environment by yourself with the setup files in the Docker folder.

First, clone this repository with submodule xgutils. xgutils contains various useful system/numpy/pytorch/3D rendering related functions that will be used by ShapeFormer.

  git clone --recursive https://github.com/QhelDIV/ShapeFormer.git

Then, create a conda environment with the yaml file.

  conda env create -f environment.yaml
  conda activate shapeformer

Next, we need to install torch_scatter through this command

  pip install torch-scatter==2.0.7 -f https://data.pyg.org/whl/torch-1.7.0+cu101.html

Demo

First, download the pretrained model from this google drive URL and extract the content to experiments/

Then run the following command to test VQDIF. The results are in experiments/demo_vqdif/results

  python -m shapeformer.trainer --opts configs/demo/demo_vqdif.yaml --gpu 0 --mode "run"

Run the following command to test ShapeFormer for shape completion. The results are in experiments/demo_shapeformer/results

  python -m shapeformer.trainer --opts configs/demo/demo_shapeformer.yaml --gpu 0 --mode "run"

Dataset

We use the dataset from IMNet, which is obtained from HSP.

The dataset we adopted is a downsampled version (64^3) from these dataset (which is 256 resolution). Please download our processed dataset from this google drive URL. And then extract the data to datasets/IMNet2_64/.

To use the full resolution dataset, please first download the original IMNet and HSP datasets, and run the make_imnet_dataset function in shapeformer/data/imnet_datasets/imnet_datasets.py

D-FAUST Human Dataset

We also provide the scripts for process the D-FAUST human shapes. First, download the official D-FAUST dataset from this link and extract to datasets/DFAUST Then, execute the following lines to generate obj files and generate sdf samples for the human meshes.

  cd shapeformer/data/dfaust_datasets/datagen
  python generate_dfaust_obj_runfile.py
  bash generate_dfaust_obj_all.sh
  python generate_dfaust_sdf_samples.py

Usage

First, train VQDIF-16 with

  python -m shapeformer.trainer --opts configs/vqdif/shapenet_res16.yaml --gpu 0

After VQDIF is trained, train ShapeFormer with

  python -m shapeformer.trainer --opts configs/shapeformer/shapenet_scale.yaml --gpu 0

For testing, you just need to append --mode test to the above commands. And if you only want to run callbacks (such as visualization/generation), set the mode to run

There is a visualization callback for both VQDIF and ShapeFormer, who will call the model to obtain 3D meshes and render them to images. The results will be save in experiments/$exp_name$/results/$callback_name$/ The callbacks will be automatically called during training and testing, so to get the generation results you just need to test the model.

ALso notice that in the configuration files batch sizes are set to very small so that the model can run on a 12GB memory GPU. You can tune it up if your GPU has a larger memory.

Multi-GPU

Notice that to use multiple GPUs, just specify the GPU ids. For example --gpu 0 1 2 4 is to use the 0th, 1st, 2nd, 4th GPU for training. Inside the program their indices will be mapped to 0 1 2 3 for simplicity.

📔 Citation

If you find our work useful for your research, please consider citing the following papers :)

@misc{yan2022shapeformer,
      title={ShapeFormer: Transformer-based Shape Completion via Sparse Representation}, 
      author={Xingguang Yan and Liqiang Lin and Niloy J. Mitra and Dani Lischinski and Danny Cohen-Or and Hui Huang},
      year={2022},
      eprint={2201.10326},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

📢: Shout-outs

The architecture of our method is inspired by ConvONet, Taming-transformers and DCTransformer. Thanks to the authors.

Also, make sure to check this amazing transformer-based image completion project(ICT)!

📧 Contact

This repo is currently maintained by Xingguang (@qheldiv) and is for academic research use only. Discussions and questions are welcome via [email protected].