Awesome
AG-Pose: Instance-Adaptive and Geometric-Aware Keypoint Learning for Category-Level 6D Object Pose Estimation
This is the official implementation of extended version of CVPR24 paper "Instance-Adaptive and Geometric-Aware Keypoint Learning for Category-Level 6D Object Pose Estimation"
[Arxiv]
The extended version primarily includes the following additions,
- Add a reconstruction network to reconstruct input point clouds using detected keypoints.
- Include experiments on HouseCat6D (CVPR 2024 Highlight) dataset.
- Include experiments using DINOv2 as image backbone.
We will soon release a preprint about the extended paper where you can find more details.
Citation
@inproceedings{lin2024instance,
title={Instance-adaptive and geometric-aware keypoint learning for category-level 6d object pose estimation},
author={Lin, Xiao and Yang, Wenfei and Gao, Yuan and Zhang, Tianzhu},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
pages={21040--21049},
year={2024}
}
Environment Settings
The code has been tested with
- python 3.9
- torch 1.12
- cuda 11.3
Some dependencies:
pip install gorilla-core==0.2.5.3
pip install opencv-python
cd model/pointnet2
python setup.py install
Data Processing
NOCS dataset
Put them under PROJ_DIR/data
and the final file structure is as follows:
data
├── camera
│ ├── train
│ ├── val
│ ├── train_list_all.txt
│ ├── train_list.txt
│ ├── val_list_all.txt
├── real
│ ├── train
│ ├── test
│ ├── train_list.txt
│ ├── train_list_all.txt
│ └── test_list_all.txt
├── segmentation_results
│ ├── CAMERA25
│ └── REAL275
├── camera_full_depths
├── gts
└── obj_models
HouseCat6D
Download and unzip the dataset from HouseCat6D and the final file structure is as follows:
HOUSECAT6D_DIR
├── scene**
├── val_scene*
├── test_scene*
└── obj_models_small_size_final
Train
Training on NOCS
python train.py --config config/REAL/camera_real.yaml
Training on HouseCat6D
python train_housecat6d.py --config config/HouseCat6D/housecat6d.yaml
Evaluate
- Evaluate on NOCS:
python test.py --config config/REAL/camera_real.yaml --test_epoch 30
- Evaluate on HouseCat6D:
python test_housecat6d.py --config config/HouseCat6D/housecat6d.yaml --test_epoch 150
Results
You can download our training logs, detailed metrics for each category and checkpoints here.
REAL275 testset:
IoU25 | IoU50 | IoU75 | 5 degree 2 cm | 5 degree 5 cm | 10 degree 2 cm | 10 degree 5 cm | |
---|---|---|---|---|---|---|---|
resnet_backbone | 84.3 | 83.8 | 77.6 | 56.2 | 62.3 | 73.4 | 81.2 |
dino_backbone | 84.3 | 84.1 | 80.1 | 57.0 | 64.6 | 75.1 | 84.7 |
CAMERA25 testset:
IoU25 | IoU50 | IoU75 | 5 degree 2 cm | 5 degree 5 cm | 10 degree 2 cm | 10 degree 5 cm | |
---|---|---|---|---|---|---|---|
resnet_backbone | 94.7 | 94.1 | 91.7 | 77.1 | 82.0 | 85.5 | 91.6 |
dino_backbone | 94.7 | 94.2 | 92.5 | 79.5 | 83.7 | 87.1 | 92.6 |
HouseCat6D testset:
IoU25 | IoU50 | IoU75 | 5 degree 2 cm | 5 degree 5 cm | 10 degree 2 cm | 10 degree 5 cm | |
---|---|---|---|---|---|---|---|
resnet_backbone | 82.4 | 66.0 | 40.5 | 11.5 | 12.6 | 37.4 | 42.5 |
dino_backbone | 88.1 | 76.9 | 53.0 | 21.3 | 22.1 | 51.3 | 54.3 |
Visualization
For visualization, please run
python visualize.py --config config/REAL/camera_real.yaml --test_epoch 30
Acknowledgements
Our implementation leverages the code from these works:
We appreciate their generous sharing.
License
Our code is released under MIT License (see LICENSE file for details).