This is the official implementation of our work PnP-Transformer. PnP-Transformer takes 2D–3D correspondences and predicts the 6D object pose, while learning a latent representation of the underlying transformation through its joint-embedding architecture. We combine PnP-Transformer with a patch classification and offset regression module to establish the 2D–3D correspondences. Our method achieves a 41.2% mean AR score across seven core BOP benchmark datasets at an inference latency of 0.46 s/obj, outperforming prior baselines such as RayPose (39.1%) and FoundPose (37.3%), while placing behind only Co-op (58.4%). When our correspondence module is combined with EPnP-RANSAC instead of PnP-Transformer, it achieves a 52.2% mean AR score at the cost of an additional 60ms per object (0.52 s/obj), ranking second overall among zero-shot coarse pose estimation methods.

Setup

Clone our repository and install the dependencies:

git clone https://github.com/ebylmz/pnp-transformer.git
cd pnp-transformer
conda env create -f environment.yaml
conda activate pnp_transformer

Install CroCo:

mkddir external
cd external/
git clone https://github.com/naver/croco.git

Compile optimized CUDA kernels for RoPE (Rotary Positional Embeddings):

cd models/curope/
python setup.py build_ext --inplace
cd ../../

Download pretrained CroCo-v2 model (external/croco/pretrained_models):

mkdir -p pretrained_models/
wget https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo_V2_ViTLarge_BaseDecoder.pth -P pretrained_models/

Using PnP-Transformer for Inference

1. Download Datasets

Download BOP datasets and CNOS detections. Your data directory should be organized in the following structure:

data/
│   bop_datasets/          
│   ├── lmo/                    # Dataset directory for LM-O
│   │   ├── models/             # 3D models of the objects
│   │   ├── models_eval/        # Simplified models for evaluation
│   │   ├── test/               # Test images and annotations
│   │   ├── camera.json
│   │   ├── dataset_info.md
│   │   ├── test_targets_bop19.json
│   │   └── ...
│   ├── tless/
│   ├── ...
├── default_detections/ 
│   ├── cnos-fastsam_lmo_test_*.json
│   ├── cnos-fastsam_tless_test_*.json
│   └── ...
├── google_scanned_objects/     # Object models (training)
├── MegaPose-GSO/               # MegaPose dataset (training)
├── onboarding/                 # RBGD template images and features
│    ├── lmo/
│    ├── tless/ 
│    └── ...
└── ...

2. Model Onboarding

Generate template images and extract features for 3D object models. PnP-Transformer can work with any backbone giving 2D-3D correspondence fields. We experimented with DINOv2, DINOv3 and CroCo-v2 and obtained best results with CroCo-v2. Here the commands to onboard models with CroCo-v2:

python scripts/onboard_bop_models.py \
    --dataset-name lmo \
    --bop-dir data/bop_datasets \
    --encoder croco \
    --checkpoint external/croco/pretrained_models/CroCo_V2_ViTLarge_BaseDecoder.pth \
    --num-subdiv 0 \
    --output-dir data/onboarding/lmo_crocov2

To be able to use DINOv3, make sure you have access to the model weights, and then authenticate via hf auth login to automatically download the model weights.

Onboarding with DINOv3:

python scripts/onboard_bop_models.py \
    --dataset-name lmo \
    --bop-dir data/bop_datasets \
    --encoder dinov3 \
    --dino-model facebook/dinov3-vitl16-pretrain-lvd1689m \
    --extract-from-layer -1 \
    --num-subdiv 0 \
    --output-dir data/onboarding/lmo_dinov3

3. Inference

Example inference on LMO dataset using default CNOS detections:

python scripts/infer_bop.py \
    --dataset-root data/bop_datasets \
    --dataset-name lmo \
    --templates-root data/onboarding \
    --detections-path data/default_detections/cnos-fastsam_lmo-test_3cb298ea-e2eb-4713-ae9e-5a7134c5da0f.json \
    --checkpoint pretrained_models/pnp_transformer_pretrained.pt \
    --pose-solver transformer 

You can also use EPnP-RANSAC for pose recovery instead of PnP-Transformer by passing 'pose-solver' as 'ransac'.

4. Evaluation on BOP datasets using bop_toolkit

First, create an 'env.sh' file with the following paths:

export REPO_PATH=/path/to/pnp-transformer
export BOP_PATH=/path/to/bop_datasets
export PYTHONPATH=$REPO_PATH:$REPO_PATH/external/bop_toolkit

Then source it:

source env.sh

Run the following script to evaluate the model predictions done on LMO dataset:

cd external/bop_toolkit
cp ../../results/lmo/pnp_transformer/pnptransformer_lmo-test.csv results/
python scripts/eval_bop19_pose.py --result_filenames pnptransformer_lmo-test.csv

Optionally, visualize the estimations using the following script:

python scripts/vis_poses.py est --result_filename pnptransformer_lmo-test.csv

Training on MegaPose Dataset

1. Downloading MegaPose-GSO Dataset

Follow the instructions from MegaPose and download the preprocessed 3D objects.

Download the training shards from HuggingFace (example for downloading first 100 shards):

python scripts/download_megapose.py \
    --dataset gso \
    --shard-start 0 \
    --shard-end 100

Generate models_info:

python scripts/generate_gso_models_info.py \
    --gso-dir data/google_scanned_objects \
    --output-file data/google_scanned_objects/models_normalized/models_info.json \
    --split normalized \
    --gso-models-json data/MegaPose-GSO/gso_models.json

Generate GSO templates:

python scripts/generate_gso_templates.py \
    --gso-dir data/google_scanned_objects \
    --gso-models-json data/MegaPose-GSO/gso_models.json \
    --num-subdiv 0 \
    --output-dir data/onboarding/gso

2. Training

Training is performed in three stages.

First, fine-tune CroCo decoder:

python scripts/train.py --config configs/training_croco.yaml --stage 1a

Then, pretrain PnP-Transformer with GT correspondences:

python scripts/train.py --config configs/training_pnp.yaml --stage 1b

Finally, post-train PnP-Transformer with estimated correspondences from stage 1a:

python scripts/train.py --config configs/training_pnp.yaml --stage 2 --croco-ckpt <path/to/finetuned_croco_decoder.pt> --pnp-ckpt <path/to/pretrained_pnp_transformer.pt> 

In any stage you can resume training from a checkpoint:

python scripts/train.py \
    --config <stage_specific_config> \
    --stage <stage> \
    --resume <path/to/checkpoint.ckpt> \
    --resume-session

Acknowledgments

We sincerely thank the authors of MegaPose for their synthetic training datasets, CroCov2 and DINOv3 for the pretrained models, and the BOP Toolkit for the evaluation tools, as well as for generously making their code and resources publicly available to the research community.

License

This project is licensed under the MIT License - see the LICENSE file for details.