FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects
Bowen Wen, Wei Yang, Jan Kautz, Stan Birchfield (NVIDIA)
CVPR 2024 |
arXiv:2312.08344 Observed crop → DINOv2 ViT-B/14 (frozen) → obs_features ─┐
├→ Concat → PoseScorer (good/bad)
Rendered crop → DINOv2 ViT-B/14 (frozen) → ren_features ─┤
└→ Concat → PoseRefiner (SE3 delta)1 import torch
2 from ergon . mlx_backend_cuda import build_cuda_model
3
4 model = build_cuda_model (
5 backbone_weights = "dinov2_vitb14_pretrain.pth" ,
6 checkpoint_path = "pytorch/ergon_v1.pth" ,
7 device = "cuda" ,
8 )
9 model . eval ( )
10
11 # Score a pose hypothesis
12 score = model . score ( observed_crop , rendered_crop ) # [B, 1] probability
13
14 # Refine a pose
15 delta = model . refine ( observed_crop , rendered_crop ) # [B, 6] (rotvec + translation)
1 import onnxruntime as ort
2 import numpy as np
3
4 sess = ort . InferenceSession ( "onnx/ergon_v1.onnx" )
5 logits , deltas = sess . run ( None , {
6 "observed" : observed_np , # [B, 3, 224, 224] float32
7 "rendered" : rendered_np , # [B, 3, 224, 224] float32
8 } )
1 # Build engine on target hardware (TRT engines are NOT portable)
2 trtexec --onnx = onnx/ergon_v1.onnx \
3 --saveEngine = tensorrt/ergon_v1_fp16.trt --fp16
├── README.md (this file)
├── pytorch/
│ ├── ergon_v1.pth (PyTorch state dict)
│ └── ergon_v1.safetensors (SafeTensors)
├── onnx/
│ └── ergon_v1.onnx (ONNX opset 17)
├── tensorrt/
│ ├── ergon_v1_fp16.trt (TensorRT FP16)
│ └── ergon_v1_fp32.trt (TensorRT FP32)
├── checkpoints/
│ └── best.pth (resume training)
├── configs/
│ └── train.yaml (training config)
└── logs/
└── training_history.json (loss curves)1 @inproceedings{wen2024foundationpose,
2 title={FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects},
3 author={Wen, Bowen and Yang, Wei and Kautz, Jan and Birchfield, Stan},
4 booktitle={CVPR},
5 year={2024}
6 }
