Addressing pose ambiguity in 6D object pose estimation from single RGB images presents a significant challenge, particularly due to object symmetries or occlusions. In response, we introduce a novel score-based diffusion method applied to the SE(3) group, marking the first application of diffusion models to SE(3) within the image domain, specifically tailored for pose estimation tasks. Extensive evaluations demonstrate the method's efficacy in handling pose ambiguity, mitigating perspective-induced ambiguity, and showcasing the robustness of our surrogate Stein score formulation on SE(3). This formulation not only improves the convergence of denoising process but also enhances computational efficiency. Thus, we pioneer a promising strategy for 6D object pose estimation.

• 2024/05/14: Code released.

Table of Contents

• Installation
  • Requirements
  • Setup
• Experiments
  • SYMSOL
  • SYMSOL-T
  • T-LESS
• Metrics
  • SYMSOL and SYMSOL-T
• Datasets
  • Customized Datasets
• Multi-GPUs Training
• Citation

Ensure your system meets the following requirements:

• Linux (only tested on Ubuntu 20.04)
• nvidia-docker
• CUDA 12.2 or higher

1. Clone this repo with the following command:

git clone [email protected]:Ending2015a/liepose-diffusion.git

2. Download datasets. This will download the TLESS dataset and VOC2012 dataset.

cd liepose-diffusion
make download

NOTE if the datasets do not download correctly, you can download them from the links provided in Datasets section.

3. Build the docker image and start the container

make build
make run
# inside the docker
cd /workspace

4. Now you are ready to run the experiments. See Experiments.

SO(3)

python3 -m liepose.exp.symsol.run

The result is located at logs/experiments/symsol-score-flat/.../inference_400000/summary.json.

SE(3)

python3 -m liepose.exp.symsolt.run "lie_type=[se3]"

R3SO(3)

python3 -m liepose.exp.symsolt.run "lie_type=[r3so3]"

The result is located at logs/experiments/symsolt-score-flat/.../inference_800000/summary.json. The ... depends on the given parameters, e.g. lie_type=se3+repr_type=tan+....

SE(3)

python3 -m liepose.exp.bop.run "lie_type=[se3]"

R3SO(3)

python3 -m liepose.exp.bop.run "lie_type=[r3so3]"

The result is located at logs/experiments/bop-tless-score-flat/.../inference_400000/summary.json. The ... depends on the given parameters, e.g. lie_type=se3+repr_type=tan+....

In summary.json, you will see the format like

{
  "final_metrics": {
    "rot": 0.007605713326483965,
    "rot(deg)": 0.4357752501964569,
    "rot_2": 99.59599999999999,
    "rot_5": 99.88,
    "rot_10": 99.94,
    "rot_id0": 0.007878238335251808,
    "rot(deg)_id0": 0.45138978958129883,
    "rot_2_id0": 99.98,
    "rot_5_id0": 100.0,
    ...
  },
  ...
}

The meaning and the shapes' ID is listed as follows:

We use the following datasets in our experiments:

*1 SYMSOL is downloaded automatically via tensorflow-datasets during the first run.

*2 We also provide the scripts for synthesizing SYMSOL-T dataset.

# You need to enable the screen/display before `make run`
export DISPLAY=:0
make run

# Make your own SYMSOL-T (25000 samples) = (5 shapes) * (5k per shape)
python3 liepose.data.symsolt.synth --path "dataset/symsolt/my-symsolt-5k" "num_samples=25000"

For more configs, see the script.

TBA.

We will announce our workflow to annotate your custom dataset in another repo Ending2015a/custom-bop-tool.

In default, single-GPU is used for training as we set the CUDA_VISIBLE_DEVICES='0' inside the Makefile. You can enable multi-GPU training simply by setting the CUDA_VISIBLE_DEVICES to multiple devices. The framework will automatically switch to the parallel training mode. For example, export CUDA_VISIBLE_DEVICES='0,1,2', will use 3 devices for training, and the batch_size is divided by 3 for each device.

@inproceedings{hsiao2024confronting,
    title={Confronting Ambiguity in 6D Object Pose Estimation via Score-Based Diffusion on SE(3)},
    author={Hsiao, Tsu-Ching and Chen, Hao-Wei and Yang, Hsuan-Kung and Lee, Chun-Yi},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
    pages={352--362},
    year={2024}
}