PVNet: Pixel-wise Voting Network for 6DoF Pose Estimation

PVNet: Pixel-wise Voting Network for 6DoF Pose Estimation
Sida Peng, Yuan Liu, Qixing Huang, Xiaowei Zhou, Hujun Bao
CVPR 2019 oral
Project Page

Introduction

Thanks Haotong Lin for providing the clean version of PVNet and reproducing the results.

Installation

1. Set up the python environment:

   python3 -m venv venv
   source venv/bin/activate
   
   # install torch 1.1 built from cuda 10.2
   pip install torch==1.1.0 torchvision==0.2.1
   
   pip install Cython==0.28.2
   sudo apt-get install libglfw3-dev libglfw3
   pip install -r requirements.txt
   

2. Compile cuda extensions under lib/csrc:

   ROOT=/path/to/clean-pvnet
   cd $ROOT/lib/csrc
   export CUDA_HOME="/usr/local/cuda-10.2"
   cd dcn_v2
   python setup.py build_ext --inplace
   cd ../ransac_voting
   python setup.py build_ext --inplace
   cd ../nn
   python setup.py build_ext --inplace
   cd ../fps
   python setup.py build_ext --inplace
   
   # If you want to use the uncertainty-driven PnP
   cd ../uncertainty_pnp
   sudo apt-get install libgoogle-glog-dev
   sudo apt-get install libsuitesparse-dev
   sudo apt-get install libatlas-base-dev
   python setup.py build_ext --inplace
   

3. Set up datasets:

   ROOT=/path/to/clean-pvnet
   cd $ROOT/data
   ln -s /path/to/custom custom
   
   

Training on the custom dataset and object

1. Use NDDS to generate a custom dataset with your own 3D-Modell (.fbx-file) containing the following:

• rgb-Image
• segmentation mask
• annotations (generated by using NVCapturableObject-tag inside UE4)

2. Process the dataset:

   python run.py --type custom
   

   Here you will be asked on which class you wish to train and on how many pictures from your dataset. The class has to be the same classname as in your annotations!

3. Train:

   python train_net.py --cfg_file configs/custom.yaml train.batch_size 4
   

   The trained models can be found in this folder: 'data/model/pvnet/custom'

4. Watch the training curve:

   tensorboard --logdir data/record/pvnet
   

5. Visualize:

   python run.py --type visualize --cfg_file configs/custom.yaml
   

6. Test:

   python run.py --type evaluate --cfg_file configs/custom.yaml
   

Citation

If you find this code useful for your research, please use the following BibTeX entry.

@inproceedings{peng2019pvnet,
  title={PVNet: Pixel-wise Voting Network for 6DoF Pose Estimation},
  author={Peng, Sida and Liu, Yuan and Huang, Qixing and Zhou, Xiaowei and Bao, Hujun},
  booktitle={CVPR},
  year={2019}
}

Acknowledgement

This work is affliated with ZJU-SenseTime Joint Lab of 3D Vision, and its intellectual property belongs to SenseTime Group Ltd.

Copyright SenseTime. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

This repository has been forked and modified to work with data from NDDS.