bruinxiong / pvn3d

Code for "A Deep Point-wise 3D Keypoints Hough Voting Network for 6DoF Pose Estimation", CVPR 2020

PVN3D

This is the source code for PVN3D: A Deep Point-wise 3D Keypoints Voting Network for 6DoF Pose Estimation, CVPR 2020. (PDF, Video).

Installation

• Install CUDA9.0/CUDA10.0
• Set up python environment from requirement.txt:

  pip3 install -r requirement.txt 

• Install tkinter through sudo apt install python3-tk
• Install python-pcl.
• Install PointNet++:

  python3 setup.py build_ext

Datasets

• LineMOD: Download the preprocessed LineMOD dataset from here (refer from DenseFusion). Unzip it and link the unzipped Linemod_preprocessed/ to pvn3d/datasets/linemod/Linemod_preprocessed:

  ln -s path_to_unzipped_Linemod_preprocessed pvn3d/dataset/linemod/

• YCB-Video: Download the YCB-Video Dataset from PoseCNN. Unzip it and link the unzippedYCB_Video_Dataset to pvn3d/datasets/ycb/YCB_Video_Dataset:

  ln -s path_to_unzipped_YCB_Video_Dataset pvn3d/datasets/ycb/

Training and evaluating

Training on the LineMOD Dataset

• First, generate synthesis data for each object using scripts from raster triangle.
• Train the model for the target object. Take object ape for example:

  cd pvn3d
  python3 -m train.train_linemod_pvn3d --cls ape

  The trained checkpoints are stored in train_log/linemod/checkpoints/{cls}/, train_log/linemod/checkpoints/ape/ in this example.

Evaluating on the LineMOD Dataset

• Start evaluation by:

  cls='ape'
  tst_mdl=train_log/linemod/checkpoints/${cls}/${cls}_pvn3d_best.pth.tar
  python3 -m train.train_linemod_pvn3d -checkpoint $tst_mdl -eval_net --test --cls $cls

  You can evaluate different checkpoint by revising tst_mdl to the path of your target model.
• We provide our pre-trained models for each object here. Download them and move them to their according folders. For example, move the ape_pvn3d_best.pth.tar to train_log/linemod/checkpoints/ape/. Then revise tst_mdl=train_log/linemod/checkpoints/ape/ape_pvn3d_best.path.tar for testing.

Training on the YCB-Video Dataset

• Preprocess the validation set to speed up training:

  cd pvn3d
  python3 -m datasets.ycb.preprocess_testset

• Start training on the YCB-Video Dataset by:

  python3 -m train.train_ycb_pvn3d

  The trained model checkpoints are stored in train_log/ycb/checkpoints/

Evaluating on the YCB-Video Dataset

• Start evaluating by:

  tst_mdl=train_log/ycb/checkpoints/pvn3d_best.pth.tar
  python3 -m train.train_ycb_pvn3d -checkpoint $tst_mdl -eval_net --test

  You can evaluate different checkpoint by revising the tst_mdl to the path of your target model.
• We provide our pre-trained models here. Download the ycb pre-trained model, move it to train_log/ycb/checkpoints/ and modify tst_mdl for testing.

Results

• Evaluation result on the LineMOD dataset:
• Evaluation result on the YCB-Video dataset:
• Visualization of some predicted poses on YCB-Video dataset:
• Joint training for distinguishing objects with similar appearance but different in size:

Citations

Please cite PVN3D if you use this repository in your publications:

@inproceedings{he2020pvn3d,
  title={PVN3D: A Deep Point-wise 3D Keypoints Voting Network for 6DoF Pose Estimation},
  author={He, Yisheng and Sun, Wei and Huang, haibin and Liu, Jianran and Fan, Haoqiang and Sun, Jian}
  booktitle={CVPR},
  year={2020}
}