/ Machine Learning

Pixel-wise Voting Network for 6DoF Pose Estimation

Pixel-wise Voting Network for 6DoF Pose Estimation

PVNet

Code for "PVNet: Pixel-wise Voting Network for 6DoF Pose Estimation" CVPR 2019 oral.

Introduction

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

The structure of this project is described in project_structure.md.

Installation

  1. Set up the python environment:
    conda create -n pvnet python=3.7
    conda activate pvnet
    
    # install torch 1.1 built from cuda 9.0
    pip install torch==1.1.0 -f https://download.pytorch.org/whl/cu90/stable
    
    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-9.0"
    cd dcn_v2
    python setup.py build_ext --inplace
    cd ../ransac_voting
    python setup.py build_ext --inplace
    cd ../fps
    python setup.py
    
  3. Set up datasets:
    ROOT=/path/to/clean-pvnet
    cd $ROOT/data
    ln -s /path/to/linemod linemod
    ln -s /path/to/linemod_orig linemod_orig
    ln -s /path/to/occlusion_linemod occlusion_linemod
    

Download datasets which are formatted for this project:

  1. linemod
  2. linemod_orig: The dataset includes the depth for each image.
  3. occlusion linemod
  4. truncation linemod: Check TRUNCATION_LINEMOD.md for the information about the Truncation LINEMOD dataset.

Testing

We provide the pretrained models of objects on Linemod, which can be found at here.

Take the testing on cat as an example.

  1. Prepare the data related to cat:
    python run.py --type linemod cls_type cat
    
  2. Download the pretrained model of cat and put it to $ROOT/data/model/pvnet/cat/199.pth.
  3. Test:
python run.py --type evaluate --cfg_file configs/linemod.yaml model cat
python run.py --type evaluate --cfg_file configs/linemod.yaml test.dataset LinemodOccTest model cat
  1. Test with icp:
python run.py --type evaluate --cfg_file configs/linemod.yaml model cat test.icp True
python run.py --type evaluate --cfg_file configs/linemod.yaml test.dataset LinemodOccTest model cat test.icp

Visualization

Take the cat as an example.

  1. Prepare the data related to cat:
    python run.py --type linemod cls_type cat
    
  2. Download the pretrained model of cat and put it to $ROOT/data/model/pvnet/cat/199.pth.
  3. Visualize:
python run.py --type visualize --cfg_file configs/linemod.yaml model cat

If setup correctly, the output will look like

cat

Training

Training

  1. Prepare the data related to cat:
    python run.py --type linemod cls_type cat
    
  2. Train:
    python train_net.py --cfg_file configs/linemod.yaml model mycat cls_type cat
    

The training parameters can be found in project_structure.md.

Tensorboard

tensorboard --logdir data/record/pvnet

If setup correctly, the output will look like

tensorboard

Training on the custom object

  1. Create a dataset using https://github.com/F2Wang/ObjectDatasetTools
  2. Organize the dataset as the following structure:
    ├── /path/to/dataset
    │   ├── model.ply
    │   ├── camera.txt
    │   ├── rgb/
    │   │   ├── 0.jpg
    │   │   ├── ...
    │   │   ├── 1234.jpg
    │   │   ├── ...
    │   ├── mask/
    │   │   ├── 0.png
    │   │   ├── ...
    │   │   ├── 1234.png
    │   │   ├── ...
    │   ├── pose/
    │   │   ├── 0.npy
    │   │   ├── ...
    │   │   ├── 1234.npy
    │   │   ├── ...
    │   │   └──
    
  3. Create a soft link pointing to the dataset:
    ln -s /path/to/custom_dataset data/custom
    
  4. Process the dataset:
    python run.py --type custom
    
  5. Train:
    python train_net.py --cfg_file configs/linemod.yaml train.dataset CustomTrain test.dataset CustomTrain model mycat train.batch_size 4
    
  6. Watch the training curve:
    tensorboard --logdir data/record/pvnet
    
  7. Visualize:
    python run.py --type visualize --cfg_file configs/linemod.yaml train.dataset CustomTrain test.dataset CustomTrain model mycat
    
  8. Test:
    python run.py --type evaluate --cfg_file configs/linemod.yaml train.dataset CustomTrain test.dataset CustomTrain model mycat
    

An example dataset can be downloaded at here.

GitHub

Comments