OnePose++：无需关键点和CAD模型的一次性物体姿态估计

项目主页 | 论文

<br/>

OnePose++：无需关键点和CAD模型的一次性物体姿态估计
何星仪^*，孙嘉铭^*，王宇昂，黄荻，鲍虎军，周晓巍
NeurIPS 2022

待办事项

• 训练、推理和演示代码。
• 在OnePose数据集和提出的OnePose_LowTexture数据集上重现评估结果的流程。
• 使用多个GPU进行并行重建和多个物体的评估。
• OnePose Cap应用：我们正在准备将数据采集应用发布到App Store（仅限iOS），敬请关注。

安装

conda env create -f environment.yaml
conda activate oneposeplus

本项目使用了LoFTR和DeepLM。感谢他们的出色工作，我们感谢他们对社区的贡献。请按照他们的安装说明和许可证进行操作：

git submodule update --init --recursive

# 安装DeepLM
cd submodules/DeepLM
sh example.sh
cp ${REPO_ROOT}/backup/deeplm_init_backup.py ${REPO_ROOT}/submodules/DeepLM/__init__.py

请注意，高效优化器DeepLM用于我们的SfM细化阶段。如果您在安装时遇到困难，请不要担心。您仍然可以通过使用我们的一阶优化器运行代码，只是速度会稍慢一些。

本项目还使用了COLMAP进行运动恢复结构。请参阅官方说明进行安装。

下载预训练模型，包括我们的2D-3D匹配和LoFTR模型。然后将它们移动到${REPO_ROOT}/weights。

[可选] 您可以选择尝试我们基于Web的3D可视化工具Wis3D，以便方便地交互式可视化特征匹配和点云。我们还在Wis3D中提供了许多其他酷炫的可视化功能，欢迎尝试。

# 正在进行中，很快就会准备就绪，目前仅在test-pypi上可用。
pip install -i https://test.pypi.org/simple/ wis3d

演示

安装完成后，您可以参考此页面使用自定义数据运行演示。

训练和评估

数据集设置

1. 从这里下载OnePose数据集，从这里下载OnePose_LowTexture数据集，并将它们解压到$/your/path/to/onepose_datasets。 如果您想在LINEMOD数据集上进行评估，请从CDPN下载真实训练数据、测试数据和3D物体模型，并从这里下载YOLOv5的检测结果。然后将它们解压到$/your/path/to/onepose_datasets/LINEMOD 目录结构应按以下方式组织：

   |--- /your/path/to/datasets
   |       |--- train_data
   |       |--- val_data
   |       |--- test_data
   |       |--- lowtexture_test_data
   |       |--- LINEMOD
   |       |      |--- real_train
   |       |      |--- real_test
   |       |      |--- models
   |       |      |--- yolo_detection
   

您可以参考数据集文档获取有关OnePose_LowTexture数据集的更多信息。

2. 建立数据集符号链接

   REPO_ROOT=/path/to/OnePose_Plus_Plus
   ln -s /your/path/to/datasets $REPO_ROOT/data/datasets

重建

对训练和测试物体都需要重建半稠密物体点云和2D-3D对应关系：

python run.py +preprocess=sfm_train_data.yaml use_local_ray=True  # 用于训练数据
python run.py +preprocess=sfm_inference_onepose_val.yaml use_local_ray=True # 用于验证数据
python run.py +preprocess=sfm_inference_onepose.yaml use_local_ray=True # 用于测试数据
python run.py +preprocess=sfm_inference_lowtexture.yaml use_local_ray=True # 用于低纹理测试数据

推理

# 评估OnePose数据集：
python inference.py +experiment=inference_onepose.yaml use_local_ray=True verbose=True

# 评估OnePose_LowTexture数据集：
python inference.py +experiment=inference_onepose_lowtexture.yaml use_local_ray=True verbose=True

请注意，我们默认在单个GPU上使用两个工作进程进行并行评估。如果您的GPU内存小于6GB，您应该添加use_local_ray=False来关闭并行化。

在LINEMOD数据集上评估

# 将LINEMOD数据集解析为OnePose数据集格式：
sh scripts/parse_linemod_objs.sh

# 在真实训练数据上重建SfM模型：
python run.py +preprocess=sfm_inference_LINEMOD.yaml use_local_ray=True

# 评估LINEMOD数据集：
python inference.py +experiment=inference_LINEMOD.yaml use_local_ray=True verbose=True

训练

1. 准备真实标注。合并训练/验证数据的标注：

   python merge.py +preprocess=merge_annotation_train.yaml
   python merge.py +preprocess=merge_annotation_val.yaml

2. 开始训练

   python train_onepose_plus.py +experiment=train.yaml exp_name=onepose_plus_train

   请注意，默认的训练配置使用8个GPU，每个GPU大约需要23GB显存。您可以在trainer.gpus中设置GPU数量或ID，并在datamodule.batch_size中减小批量大小以降低GPU显存占用。

所有模型权重将保存在${REPO_ROOT}/models/checkpoints/${exp_name}下，日志将保存在${REPO_ROOT}/logs/${exp_name}下。 您可以通过Tensorboard可视化训练过程：

tensorboard --logdir logs --bind_all --port your_port_number

引用

如果您发现这段代码对您的研究有用，请使用以下BibTeX条目。

@inproceedings{
    he2022oneposeplusplus,
    title={OnePose++: Keypoint-Free One-Shot Object Pose Estimation without {CAD} Models},
    author={Xingyi He and Jiaming Sun and Yuang Wang and Di Huang and Hujun Bao and Xiaowei Zhou},
    booktitle={Advances in Neural Information Processing Systems},
    year={2022}
}

致谢

我们的部分代码借鉴自hloc和LoFTR。感谢他们的作者的出色工作。