ShAPO:tophat:: Implicit Representations for Multi-Object Shape, Appearance and Pose Optimization

This repository is the pytorch implementation of our paper:

ShAPO: Implicit Representations for Multi-Object Shape, Appearance and Pose Optimization
Muhammad Zubair Irshad, Sergey Zakharov, Rares Ambrus, Thomas Kollar, Zsolt Kira, Adrien Gaidon
European Conference on Computer Vision (ECCV), 2022

[Project Page] [arXiv] [PDF] [Video] [Poster]

Citation

If you find this repository useful, please consider citing:

@inproceedings{irshad2022shapo,
  title={ShAPO: Implicit Representations for Multi Object Shape Appearance and Pose Optimization},
  author={Muhammad Zubair Irshad and Sergey Zakharov and Rares Ambrus and Thomas Kollar and Zsolt Kira and Adrien Gaidon},
  journal={European Conference on Computer Vision (ECCV)},
  year={2022},
  url={https://arxiv.org/abs/2207.13691},
}

@inproceedings{irshad2022centersnap,
  title={CenterSnap: Single-Shot Multi-Object 3D Shape Reconstruction and Categorical 6D Pose and Size Estimation},
  author={Muhammad Zubair Irshad and Thomas Kollar and Michael Laskey and Kevin Stone and Zsolt Kira},
  journal={IEEE International Conference on Robotics and Automation (ICRA)},
  year={2022},
  url={https://arxiv.org/abs/2203.01929},
}

💻 Environment
📊 Dataset
✨ Training and Inference
📝 FAQ (Updated)

💻 Environment

Create a python 3.8 virtual environment and install requirements:

cd $ShAPO_Repo
conda create -y --prefix ./env python=3.8
conda activate ./env/
./env/bin/python -m pip install --upgrade pip
./env/bin/python -m pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html

The code was built and tested on cuda 10.2

📊 Dataset

Download camera_train, camera_val, real_train, real_test, ground-truth annotations, camera_composed_depth, mesh models and eval_resultsprovided by NOCS and nocs preprocess data.
Also download sdf_rgb_pretrained_weights. Unzip and organize these files in $ShAPO_Repo/data as follows:

data
├── CAMERA
│   ├── train
│   └── val
├── Real
│   ├── train
│   └── test
├── camera_full_depths
│   ├── train
│   └── val
├── gts
│   ├── val
│   └── real_test
├── auto_encoder_model
│   ├── model_50_nocs.pth
├── results
│   ├── camera
│   ├── mrcnn_results
│   ├── nocs_results
│   └── real
├── sdf_rgb_pretrained
│   ├── LatentCodes
│   ├── Reconstructions
│   ├── ModelParameters
│   ├── OptimizerParameters
│   └── rgb_latent
└── obj_models
    ├── train
    ├── val
    ├── real_train
    ├── real_test
    ├── camera_train.pkl
    ├── camera_val.pkl
    ├── real_train.pkl
    ├── real_test.pkl
    └── mug_meta.pkl

Create image lists

./runner.sh prepare_data/generate_training_data.py --data_dir /home/ubuntu/shapo/data/nocs_data/

Now run distributed script to collect data locally in a few hours. The data would be saved under data/NOCS_data.

Note: The script uses multi-gpu and runs 8 workers per gpu on a 16GB GPU. Change worker_per_gpu variable depending on your GPU size.

python prepare_data/distributed_generate_data.py --data_dir /home/ubuntu/shapoplusplus/data/nocs_data --type camera_train

--type chose from 'camera_train', 'camera_val', 'real_train', 'real_val'

✨ Training and Inference

ShAPO is a two-stage process; First, a single-shot network to predict 3D shape, pose and size codes along with segmentation masks in a per-pixel manner. Second, test-time optimization of joint shape, pose and size codes given a single-view RGB-D observation of a new instance.

Train on NOCS Synthetic (requires 13GB GPU memory):

./runner.sh net_train.py @configs/net_config.txt

Note than runner.sh is equivalent to using python to run the script. Additionally it sets up the PYTHONPATH and CenterSnap Enviornment Path automatically. Also note that this part of the code is similar to CenterSnap. We predict implicit shapes as SDF MLP instead of pointclouds and additionally also predict appearance embedding and object masks in this stage.

Finetune on NOCS Real Train (Note that good results can be obtained after finetuning on the Real train set for only a few epochs i.e. 1-5):

./runner.sh net_train.py @configs/net_config_real_resume.txt --checkpoint \path\to\best\checkpoint

Inference on a NOCS Real Test Subset

Download a small Real test subset from [here],our shape and texture pretrained checkpoints from [here] and shapo pretrained checkpoints on real dataset here. Unzip and organize these files in $ShAPO_Repo/data as follows:

test_data
├── Real
│   ├── test
|   ckpts

└── sdf_rgb_pretrained
    ├── LatentCodes
    ├── LatentCodes
    ├── Reconstructions
    ├── ModelParameters
    ├── OptimizerParameters
    └── rgb_latent

Now run the inference script to visualize the single-shot predictions as follows:

bash
./runner.sh inference/inference_real.py @configs/net_config.txt --test_data_dir path_to_nocs_test_subset --checkpoint checkpoint_path_here

You should see the visualizations saved in results/ShAPO_real. Change the --ouput_path in *config.txt to save them to a different folder

Optimization

This is the core optimization script to update latent shape and appearance codes along with 6D pose and sizes to better the fit the unseen single-view RGB-D observation. For a quick run of the core optimization loop along with visualization, see this notebook here

./runner.sh opt/optimize.py @configs/net_config.txt --data_dir /path/to/test_data_dir/ --checkpoint checkpoint_path_here

📝 FAQ

Please see FAQs from CenterSnap here

Acknowledgments

This code is built upon the implementation from CenterSnap

Related Work

CenterSnap: Single-Shot Multi-Object 3D Shape Reconstruction and Categorical 6D Pose and Size Estimation, ICRA, 2022

Licenses

This repository is released under the CC BY-NC 4.0 license.

AtlasGooo2/shapo