LiZheng1997 / 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_bilibili, Video_youtube).

Installation

• Install CUDA9.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++ (refer from Pointnet2_PyTorch):

  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:

  # commands in eval_linemod.sh
  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 at onedrive link, baiduyun link (access code(提取码): 8kmp). 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.

Demo/visualizaion on the LineMOD Dataset

• After training your models or downloading the pre-trained models, you can start the demo by:

  # commands in demo_linemod.sh
  cls='ape'
  tst_mdl=train_log/linemod/checkpoints/${cls}/${cls}_pvn3d_best.pth.tar
  python3 -m demo -dataset linemod -checkpoint $tst_mdl -cls $cls

  The visualization results will be stored in train_log/linemod/eval_results/{cls}/pose_vis

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:

  # commands in eval_ycb.sh
  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 at onedrive link, baiduyun link (access code(提取码): h2i5). Download the ycb pre-trained model, move it to train_log/ycb/checkpoints/ and modify tst_mdl for testing.

Demo/visualizaion on the LineMOD Dataset

• After training your model or downloading the pre-trained model, you can start the demo by:

  # commands in demo_ycb.sh
  tst_mdl=train_log/ycb/checkpoints/pvn3d_best.pth.tar
  python3 -m demo -checkpoint $tst_mdl -dataset ycb

  The visualization results will be stored in train_log/ycb/eval_results/pose_vis

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{he_2020_CVPR,
  author={He, Yisheng and Sun, Wei and Huang, haibin and Liu, Jianran and Fan, Haoqiang and Sun, Jian}
  title={PVN3D: A Deep Point-wise 3D Keypoints Voting Network for 6DoF Pose Estimation},
  booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  month = {June},
  year={2020}
}