Pseudo Flow Consistency for Self-Supervised 6D Object Pose Estimation (ICCV 2023)

Yang Hai, Rui Song, Jiaojiao Li, David Ferstl, Yinlin Hu

Paper | Poster | Video

Introduction

Most self-supervised 6D object pose estimation methods can only work with additional depth information or rely on the accurate annotation of 2D segmentation masks, limiting their application range. In this paper, we propose a 6D ob- ject pose estimation method that can be trained with pure RGB images without any auxiliary information.

Installation

1. Install necessary packages by pip install -r requirements.txt
2. Install Pytorch3D by building this code repo from source.
3. Install bop_toolkit(Optional).

Dataset Preparation

1. Download BOP YCB-V Dataset, and place them under data directory.
2. Download the image lists and other data related to YCB-V, which will be used in this code, from here.
3. Download the detected bounding boxes by RADet, from here.

Training

1. Train the pose initialization network, which is WDR-Pose equipped with RADet pre-processing, on synthetic images.

python train.py --config configs/estimator/extended_wdr.py --mode estimator

2. Train the optical flow network on synthetic images.

python train.py --config configs/flow_refine/raft_flow_mask.py --mode refiner

3. Infer the initial poses of un-annotated real training images.

python test.py --config configs/estimator/extended_wdr.py --mode estimator --checkpoint work_dirs/wdr_ycbv_pbr/latest.pth --format-only --save-dir data/initial_poses/extended_wdr/ycbv_pbr_train

4. Train the optical flow network on un-annotated real images, which will load the pretrained weights and use the above initial poses for training.

python train.py --config configs/flow_refine/pfc_raft_flow_mask.py --mode refiner

Notes: To reproduce the results in our paper, the first three steps can be skipped by accessing the pretrained model weights from here and initial pose of unlabeled real images from here.

Testing

1. Infer the initial poses on testing images, and save them.

python test.py --config configs/estimator/extended_wdr.py --checkpoint work_dirs/wdr_ycbv_pbr/latest.pth --format-only --save-dir data/initial_poses/extended_wdr/ycbv_pbr_test

2. Run the trained optical flow network to refine the above initial poses.

python test.py --config configs/flow_refine/pfc_raft_flow_mask.py --checkpoint work_dirs/pfc_real_selfsup/latest.pth --eval

Notes: The trained model on YCB-V can be found here and initial pose for testing images can be found here.

Evaluation under BOP setting(optional)

1. Save the results.

python test.py --config configs/flow_refine/pfc_raft_flow_mask.py --checkpoint work_dirs/pfc_real_selfsup/latest.pth --format-only --save-dir work_dirs/pfc_real_seflsup/results

2. Convert to BOP format.

python tools/convert_to_bop19.py work_dirs/pfc_real_seflsup/results data/ycbv/test_targets_bop19.json work_dirs/pfc_real_seflsup/results_bop19.json

3. Follow the bop_toolkit instruction for evaluation.

Citation

If you find this project is helpful, please cite:

@inproceedings{yang2023pseudoflow,
  title={Pseudo Flow Consistency for Self-Supervised 6D Object Pose Estimation},
  author={Yang Hai and Rui Song and Jiaojiao Li and David Ferstl and Yinlin Hu},
  booktitle={ICCV},
  year={2023}
}
@inproceedings{yang2023radet,
  title={Rigidity-Aware Detection for 6D Object Pose Estimation},
  author={Yang Hai and Rui Song and Jiaojiao Li and Mathieu Salzmann and Yinlin Hu},
  booktitle={CVPR},
  year={2023}
}
@inproceedings{yang2023scflow,
  title={Shape-Constraint Recurrent Flow for 6D Object Pose Estimation},
  author={Yang Hai and Rui Song and Jiaojiao Li and Yinlin Hu},
  booktitle={CVPR},
  year={2023}
}
@inproceedings{yinlin2022wdr,
  title={Wide-Depth-Range 6D Object Pose Estimation in Space},
  author={Yinlin Hu, Se ́bastien Speierer, Wenzel Jakob, Pascal Fua, Mathieu Salzmann},
  booktitle={CVPR},
  year={2022}
}

Notes

Some parts of the code are based on WDR-Pose and SCFlow.