PoseRBPF

PoseRBPF: A Rao-Blackwellized Particle Filter for 6D Object Pose Tracking

Citing PoseRBPF

If you find the PoseRBPF code useful, please consider citing:

@inproceedings{deng2019pose,
author    = {Xinke Deng and Arsalan Mousavian and Yu Xiang and Fei Xia and Timothy Bretl and Dieter Fox},
title     = {PoseRBPF: A Rao-Blackwellized Particle Filter for 6D Object Pose Tracking},
booktitle = {Robotics: Science and Systems (RSS)},
year      = {2019}
}
@inproceedings{deng2020self,
author    = {Xinke Deng and Yu Xiang and Arsalan Mousavian and Clemens Eppner and Timothy Bretl and Dieter Fox},
title     = {Self-supervised 6D Object Pose Estimation for Robot Manipulation},
booktitle = {International Conference on Robotics and Automation (ICRA)},
year      = {2020}
}

Installation

git clone https://github.com/NVlabs/PoseRBPF.git --recursive

Install dependencies:

  • install anaconda according to the official website.
  • create the virtual env with pose_rbpf_env.yml:
conda env create -f pose_rbpf_env.yml
conda activate pose_rbpf_env
  • compile the YCB Renderer according to the instruction.
  • compile the utility functions with:
sh build.sh

Download

Downolad files as needed. Extract CAD models under the cad_models directory, and extract model weights under the checkpoints directory.

A quick demo on the YCB Video Dataset

demo

  • The demo shows tracking 003_cracker_box on YCB Video Dataset.
  • Run script download_demo.sh to download checkpoint (434 MB), CAD models (743 MB), 2D detections (13 MB), and necessary data (3 GB) for the demo:
./scripts/download_demo.sh
  • Then you should have files organized like:
├── ...
├── PoseRBPF
|   |── cad_models
|   |   |── ycb_models
|   |   └── ...
|   |── checkpoints
|   |   |── ycb_ckpts_roi_rgbd
|   |   |── ycb_codebooks_roi_rgbd
|   |   |── ycb_configs_roi_rgbd
|   |   └── ... 
|   |── detections
|   |   |── posecnn_detections
|   |   |── tless_retina_detections 
|   |── config                      # configuration files for training and DPF
|   |── networks                    # auto-encoder networks
|   |── pose_rbpf                   # particle filters
|   └── ...
|── YCB_Video_Dataset               # to store ycb data
|   |── cameras
|   |── data 
|   |── image_sets 
|   |── keyframes 
|   |── poses
|   └── ...
└── ...
  • Run demo with 003_cracker_box. The results will be stored in ./results/
./scripts/run_demo.sh

Online Real-world Pose Estimation using ROS

ros_tracking

  • Due to the incompatibility between ROS Kinetic and Python 3, the ROS node only runs with Python 2.7. We first create the virtual env with pose_rbpf_env_py2.yml:
conda env create -f pose_rbpf_env_py2.yml
conda activate pose_rbpf_env_py2
  • compile the YCB Renderer according to the instruction.
  • compile the utility functions with:
sh build.sh
  • Make sure you can run the demo above first.
  • Install ROS if it's not there:
sudo sh -c 'echo "deb http://packages.ros.org/ros/ubuntu $(lsb_release -sc) main" > /etc/apt/sources.list.d/ros-latest.list'
sudo apt-key adv --keyserver 'hkp://keyserver.ubuntu.com:80' --recv-key C1CF6E31E6BADE8868B172B4F42ED6FBAB17C654
sudo apt-get update
sudo apt-get install ros-kinetic-desktop-full
  • Update python packages:
conda install -c auto catkin_pkg
pip install -U rosdep rosinstall_generator wstool rosinstall six vcstools
pip install msgpack
pip install empy
  • Source ROS (every time before launching the node):
source /opt/ros/kinetic/setup.bash
  • Initialze rosdep:
sudo rosdep init
rosdep update

Single object tracking demo:

  • Download demo rosbag:
./scripts/download_ros_demo.sh
  • Run PoseCNN node (with roscore running in another terminal, download PoseCNN weights first):
./scripts/run_ros_demo_posecnn.sh
  • Run PoseRBPF node for RGB-D tracking (with roscore running in another terminal):
./scripts/run_ros_demo.sh
  • (Optional) For RGB tracking run this command instead:
./scripts/run_ros_demo_rgb.sh
  • Run RVIZ in the PoseRBPF directory:
rosrun rviz rviz -d ./ros/tracking.rviz
  • Once you see *** PoseRBPF Ready ... in the PoseRBPF terminal, run rosbag in another terminal, then you should be able to see the tracking demo:
rosbag play ./ros_data/demo_single.bag

Multiple object tracking demo:

  • Download demo rosbag:
./scripts/download_ros_demo_multiple.sh
  • Run PoseCNN node (with roscore running in another terminal, download PoseCNN weights first):
./scripts/run_ros_demo_posecnn.sh
  • Run PoseRBPF node with self-supervised trained RGB Auto-encoder weights:
./scripts/run_ros_demo_rgb_multiple_ssv.sh
  • (Optional) Run PoseRBPF node with RGB-D Auto-encoder weights instead:
./scripts/run_ros_demo_multiple.sh
  • (Optional) Run PoseRBPF node with RGB Auto-encoder weights instead:
./scripts/run_ros_demo_rgb_multiple.sh
  • Run RVIZ in the PoseRBPF directory:
rosrun rviz rviz -d ./ros/tracking.rviz
  • Once you see *** PoseRBPF Ready ... in the PoseRBPF terminal, run rosbag in another terminal, then you should be able to see the tracking demo:
rosbag play ./ros_data/demo_multiple.bag

Note that PoseRBPF takes certain time to initialize each object before tracking. You can pause the ROS bag by pressing space for initialization, and then press space again to resume tracking.

Testing on the YCB Video Dataset

  • Download checkpoints from the google drive folder (ycb_rgbd_full.tar.gz or ycb_rgb_full.tar.gz) and unzip to the checkpoint directory.
  • Download all the data in the YCB Video Dataset so the ../YCB_Video_Dataset/data folder contains all the sequences.
  • Run RGB-D tracking (use 002_master_chef_can as an example here):
sh scripts/test_ycb_rgbd/val_ycb_002_rgbd.sh 0 1
  • Run RGB tracking (use 002_master_chef_can as an example here):
sh scripts/test_ycb_rgb/val_ycb_002_rgb.sh 0 1

Testing on the T-LESS Dataset

  • Download checkpoints from the google drive folder (tless_rgbd_full.tar.gz or tless_rgb_full.tar.gz) and unzip to the checkpoint directory.
  • Download all the data in the T-LESS Dataset so the ../TLess/ folder contains all the sequences.
  • Download all the models for T-LESS objects from the google drive folder.
  • Then you should have files organized like:
├── ...
├── PoseRBPF
|   |── cad_models
|   |   |── ycb_models
|   |   |── tless_models
|   |   └── ...
|   |── checkpoints
|   |   |── tless_ckpts_roi_rgbd
|   |   |── tless_codebooks_roi_rgbd
|   |   |── tless_configs_roi_rgbd
|   |   └── ... 
|   |── detections
|   |   |── posecnn_detections
|   |   |── tless_retina_detections 
|   |── config                      # configuration files for training and DPF
|   |── networks                    # auto-encoder networks
|   |── pose_rbpf                   # particle filters
|   └── ...
|── YCB_Video_Dataset               # to store ycb data
|   |── cameras  
|   |── data 
|   |── image_sets 
|   |── keyframes 
|   |── poses               
|   └── ...   
|── TLess               # to store tless data
|   |── t-less_v2 
|── tless_ckpts_roi_rgbd
|   |   |── test_primesense
|   |   └── ... 
|   └── ...        
└── ...
  • Run RGB-D tracking (use obj_01 as an example here):
sh scripts/test_tless_rgbd/val_tless_01_rgbd.sh 0 1
  • Run RGB tracking (use obj_01 as an example here):
sh scripts/test_tless_rgb/val_tless_01_rgb.sh 0 1

Testing on the DexYCB Dataset

  • Download checkpoints from the google drive folder (ycb_rgbd_full.tar.gz or ycb_rgb_full.tar.gz) and unzip to the checkpoint directory.

  • Download the DexYCB dataset from here.

  • Download PoseCNN results on the DexYCB dataset from here.

  • Create a symlink for the DexYCB dataset and the PoseCNN results

    cd $ROOT/data/DEX_YCB
ln -s $dex_ycb_data data
ln -s $results_posecnn_data results_posecnn

  • Install PyTorch PoseCNN layers according to the instructions here.

  • Run RGB-D tracking:

    ./scripts/test_dex_rgbd/dex_ycb_test_rgbd_s0.sh $GPU_ID

  • Run RGB tracking:

    ./scripts/test_dex_rgb/dex_ycb_test_rgb_s0.sh $GPU_ID

Training

  • Download microsoft coco dataset 2017 val images from here for data augmentation.
  • Store the folder val2017 in ../coco/
  • Run training example for 002_master_chef_can in the YCB objects. The training should be able to run on one single NVIDIA TITAN Xp GPU:
sh scripts/train_ycb_rgbd/train_script_ycb_002.sh

Acknowledgements

We have referred to part of the RoI align code from maskrcnn-benchmark.

GitHub

https://github.com/NVlabs/PoseRBPF