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

• 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

• Download and preprocess the dataset following DPDN
• Download and unzip the segmentation results here

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

• 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:

CAMERA25 testset:

HouseCat6D testset:

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:

• NOCS
• SPD
• DualPoseNet
• DPDN
• VI-Net
• HouseCat6D Toolbox

We appreciate their generous sharing.

License

Our code is released under MIT License (see LICENSE file for details).

Contact

llinxiao@mail.ustc.edu.cn