This repository contains the code for the paper Category Level Object Pose Estimation via Neural Analysis-by-Synthesis.
This code has been tested in Ubuntu 18.04 and Python 3.7.7.
conda create -n neural_object_fitting python=3.7
conda activate neural_object_fitting
conda install pytorch=1.1.0 torchvision=0.3.0 cudatoolkit=10.0 -c pytorch
pip install -r requirements.txt
Run the following command to download the datasets and checkpoints. This script downloads the real test dataset from NOCS, our synthetic training set, and our pre-trained models. Our method uses 2D detections and segmentations predicted by NOCS. This script also downloads the pre-computed results. You can also obtain these predictions by running NOCS following the instructions in the original repo
sh prepare_datasets.sh
Run the following command to estimate poses.
python eval.py --config configs/eval.yaml
Note: a full evaluation can take ~2 days on a single 1080Ti GPU. You can split the evaluation into multiple batches which can then be run in parallel on a cluster. For example, the following command splits the job into two parts, and each agent takes care of one.
python eval.py --config configs/eval.yaml --num_agent 2 --id_agent 0
python eval.py --config configs/eval.yaml --num_agent 2 --id_agent 1
You can also run the evaluation on a subset of the data first for a sanity check. This can be done by specifying --skip option, e.g. with --skip 10 only 1/10 data will be evaluated.
python eval.py --config configs/eval.yaml --skip 10
To visualize the fitting procedure, add --vis to the command.
python eval.py --config configs/eval.yaml --vis
Run the following command to evaluate the estimated poses and draw the plot.
python nocs/eval.py
This should produce similar results to this one. Note that there might be a small variance due to the randomness of initialization.
Run the following command to train the model for laptop category. Modify --category to train models for other categories. Each of the provided checkpoints was trained for 24h on a single 1080Ti GPU.
python train.py --config configs/config.yaml --category laptop
Note: we use Weight&Bias for experiment logging which requires registration (for free). To avoid registration, you can deactivate the usage of Weight&Bias by removing wandb from visualizers in configs/config.yaml.
If you find this repository useful, please consider citing our paper.
@article{chen2020category,
title={Category Level Object Pose Estimation via Neural Analysis-by-Synthesis},
author={Chen, Xu and Dong, Zijian, and Song, Jie and Geiger, Andreas and Hilliges, Otmar},
year= {2020},
booktitle = {European Conference on Computer Vision (ECCV)},
}
This code is based on pytorch-CycleGAN-and-pix2pix.