Home

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,

  1. Add a reconstruction network to reconstruct input point clouds using detected keypoints.
  2. Include experiments on HouseCat6D (CVPR 2024 Highlight) dataset.
  3. 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

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/dataand 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

python test.py --config config/REAL/camera_real.yaml --test_epoch 30
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:

IoU25IoU50IoU755 degree 2 cm5 degree 5 cm10 degree 2 cm10 degree 5 cm
resnet_backbone84.383.877.656.262.373.481.2
dino_backbone84.384.180.157.064.675.184.7

CAMERA25 testset:

IoU25IoU50IoU755 degree 2 cm5 degree 5 cm10 degree 2 cm10 degree 5 cm
resnet_backbone94.794.191.777.182.085.591.6
dino_backbone94.794.292.579.583.787.192.6

HouseCat6D testset:

IoU25IoU50IoU755 degree 2 cm5 degree 5 cm10 degree 2 cm10 degree 5 cm
resnet_backbone82.466.040.511.512.637.442.5
dino_backbone88.176.953.021.322.151.354.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).

Contact

llinxiao@mail.ustc.edu.cn