Home

Awesome

DGECN

This repo provides the PyTorch implementation of the work:

Tuo Cao, Fei Luo, Yanping Fu, Wenxiao Zhang, Shengjie Zheng, Chunxia Xiao. DGECN: A Depth-Guided Edge Convolutional Network for End-to-End 6D Pose Estimation. In CVPR 2022. Project page

Please note that this repo only contains test code without KFA and DRN, and test with RANSAC/PnP. But we also provide our DG-PnP, we may provide the full code in Huawei Ascend after the conference.

Overview

<p align="center"> <img src='pic/overview.png' width='800'> <p>

Requirements

Datasets and Pretrained model

Download the YCB-V from here and extract to ./data.

Please also download the pretrained model from here (BaiduNetDisk, OneDrive, password: gk8y).

Evaluation on YCB-V

python test.py --use_gpu --filelist=FILELIST --out_dir=./Result --test_mode=YCB-Video --model_path=MODEL_PATH

Example:

python test.py --use_gpu --filelist=ycb-video-testlist.txt --out_dir=./Result --test_mode=YCB-Video --model_path=dgecn.pth

Visualize depth predictions

python depth_vis.py --image_path=IMAGE_PATH --model_path=MODEL_PATH --ext==FILE_EXT

Example:

python depth_vis.py --image_path=./assert/test_0.png --model_path=dgecn.pth --ext=png

This pretrained model is trained on video 0000 ~ 0010， the predictions on other videos may not good. We will release our pretrained model on all videos in the full code.

<p align="center"> <img src="assert/test_0.png" width="400"/><img src="assert/test_0_disp.jpeg" width="400"/> </p>

Citation

If you find this useful in your research, please consider citing our paper.

@InProceedings{Cao_2022_CVPR,
    author    = {Cao, Tuo and Luo, Fei and Fu, Yanping and Zhang, Wenxiao and Zheng, Shengjie and Xiao, Chunxia},
    title     = {DGECN: A Depth-Guided Edge Convolutional Network for End-to-End 6D Pose Estimation},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {3783-3792}
}

References