RBP-Pose

Pytorch implementation of RBP-Pose: Residual Bounding Box Projection for Category-Level Pose Estimation.

Required environment

• Ubuntu 18.04
• Python 3.8
• Pytorch 1.10.1
• CUDA 11.3.

Installing

• Install the main requirements in 'requirement.txt'.
• Install Detectron2.

Data Preparation

To generate your own dataset, use the data preprocess code provided in this git. Download the detection results in this link.

Trained model

Trained model is available here.

Training

Please note, some details are changed from the original paper for more efficient training.

Specify the dataset directory and run the following command.

python -m engine.train --data_dir YOUR_DATA_DIR --model_save SAVE_DIR --training_stage shape_prior_only # first stage
python -m engine.train --data_dir YOUR_DATA_DIR --model_save SAVE_DIR --resume 1 --resume_model MODEL_PATH--training_stage prior+recon+novote # second stage

Detailed configurations are in 'config/config.py'.

Evaluation

python -m evaluation.evaluate --data_dir YOUR_DATA_DIR --detection_dir DETECTION_DIR --resume 1 --resume_model MODEL_PATH --model_save SAVE_DIR

Acknowledgment

Our implementation leverages the code from 3dgcn, FS-Net, DualPoseNet, SPD.