CenterSnap：单次多目标三维形状重建和类别6D姿态与尺寸估计

本仓库是我们论文的PyTorch实现：

CenterSnap：单次多目标三维形状重建和类别6D姿态与尺寸估计
Muhammad Zubair Irshad，Thomas Kollar，Michael Laskey，Kevin Stone，Zsolt Kira
2022年国际机器人与自动化会议（ICRA）

[项目页面] [arXiv] [PDF] [视频] [海报]

后续ECCV'22工作：

ShAPO：多目标形状、外观和姿态优化的隐式表示
Muhammad Zubair Irshad，Sergey Zakharov，Rares Ambrus，Thomas Kollar，Zsolt Kira，Adrien Gaidon
2022年欧洲计算机视觉会议（ECCV）

[项目页面] [arXiv] [PDF] [视频] [海报]

引用

如果您觉得本仓库有用，请考虑引用：

@inproceedings{irshad2022centersnap,
     title = {CenterSnap: 单次多目标3D形状重建和类别6D姿态与尺寸估计},
     author = {Muhammad Zubair Irshad and Thomas Kollar and Michael Laskey and Kevin Stone and Zsolt Kira},
     journal = {IEEE国际机器人与自动化会议(ICRA)},
     year = {2022}
     }


@inproceedings{irshad2022shapo,
     title = {ShAPO: 多目标形状外观和姿态优化的隐式表示},
     author = {Muhammad Zubair Irshad and Sergey Zakharov and Rares Ambrus and Thomas Kollar and Zsolt Kira and Adrien Gaidon},
     journal = {欧洲计算机视觉会议(ECCV)},
     year = {2022}
     }

💻 环境

创建Python 3.8虚拟环境并安装依赖项：

cd $CenterSnap_Repo
conda create -y --prefix ./env python=3.8
conda activate ./env/
./env/bin/python -m pip install --upgrade pip
./env/bin/python -m pip install -r requirements.txt

根据您的CUDA版本安装torch==1.7.1 torchvision==0.8.2。该代码在cuda 10.2上构建和测试。在cuda 10.2上安装torch的示例命令如下：

pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2

📊 数据集

最新更新：如果您想从头开始收集自己的数据，请查看我们新的ECCV'22工作ShAPO的分布式脚本。该分布式脚本收集的数据格式与CenterSnap所需的格式相同，只需进行一些小的修改，详见该仓库。

下载预处理数据集

我们建议下载预处理的数据集来训练和评估CenterSnap模型。下载并解压合成数据集（868GB）和真实数据集（70GB）。这些文件包含复现我们结果所需的所有训练和验证数据。

cd $CenterSnap_REPO/data
wget https://tri-robotics-public.s3.amazonaws.com/centersnap/CAMERA.tar.gz
tar -xzvf CAMERA.tar.gz

wget https://tri-robotics-public.s3.amazonaws.com/centersnap/Real.tar.gz
tar -xzvf Real.tar.gz

数据目录结构应如下所示：

data
├── CAMERA
│   ├── train
│   └── val_subset
├── Real
│   ├── train
└── └── test

要准备自己的数据集，我们在prepare_data下提供了额外的脚本。

✨ 训练和推理

在NOCS合成数据集上训练（需要13GB GPU内存）：

./runner.sh net_train.py @configs/net_config.txt

注意，runner.sh等同于使用python运行脚本。此外，它还自动设置PYTHONPATH和CenterSnap环境路径。

在NOCS真实训练集上微调（请注意，在真实训练集上微调几个epoch（即1-5个）就可以获得不错的结果）：

./runner.sh net_train.py @configs/net_config_real_resume.txt --checkpoint \path\to\best\checkpoint

在NOCS真实测试子集上进行推理

从[这里]下载一个小型NOCS Real子集

./runner.sh inference/inference_real.py @configs/net_config.txt --data_dir nocs_test_subset路径 --checkpoint checkpoint路径

您应该会看到可视化结果保存在results/CenterSnap中。更改*config.txt中的--ouput_path可将它们保存到不同的文件夹

可选（形状自编码器预训练）

我们提供了预训练的形状自编码器模型，用于数据收集和推理。虽然我们的代码库不需要单独训练形状自编码器，但如果您想这样做，我们在external/shape_pretraining下提供了额外的脚本

📝 常见问题

1. 我在自定义相机图像（如Realsense、OAK-D或其他）上无法获得良好的性能。

答：由于网络仅在真实世界NOCS数据上进行了微调，目前预训练网络仅在以下相机设置上给出良好的3D预测。要在您自己的相机参数上获得良好的预测，请确保在合成数据集上预训练后，使用您自己的小型子集微调网络。我们在这里提供了数据准备脚本。

2. 如何在HSR机器人相机上生成良好的零样本结果：

答：请参阅上面FAQ1的答案以获得最佳结果。我们采用的另一种解决方案是对HSR机器人相机或任何其他自定义相机输出的rgb-d观察进行变形，使其与NOCS真实相机（我们在其上微调模型）的内参匹配。这样，只需在NOCS真实数据集上微调，就可以获得不错的结果。请参阅此答案和相应的gist以获取代码。

3. 运行colab时出现no cuda GPUs available。

答：请确保按照以下说明在colab中激活GPU：

确保您在Runtime-> Change runtime type下启用了GPU！

4. 我遇到raise RuntimeError('received %d items of ancdata' % RuntimeError: received 0 items of ancdata

答：通过ulimit -n 2048将ulimit增加到2048或8096

5. 我遇到 RuntimeError: CUDA error: no kernel image is available for execution on the device或You requested GPUs: [0] But your machine only has: []

答：检查您的pytorch安装是否与cuda安装匹配。尝试以下操作：

安装cuda 10.2并运行requirements.txt中的相同脚本
安装相关的pytorch cuda版本，即更改requirements.txt中的这一行

torch==1.7.1
torchvision==0.8.2

6. 我在wandb中看到零值验证指标

答：确保对指标进行阈值处理。由于pytorch lightning的第一次验证检查指标较高，其他指标似乎都为零。请在wandb中手动设置阈值以删除异常指标，以查看实际指标。

后续工作

我们的ECCV后续工作：
- ShAPO：多目标形状外观和姿态优化的隐式表示，ECCV，2022
其他后续工作（感谢作者们的出色工作）：
- 铰接物体：CARTO：类别和关节无关的铰接物体重建，CVPR，2023
- 抓取：实时同步多目标3D形状重建、6自由度姿态估计和密集抓取预测
- 更多抓取 + 真实世界实验：CenterGrasp：用于同时进行形状重建和6自由度抓取估计的物体感知隐式表示学习

致谢

本代码基于SimNet的实现构建

许可证

源代码根据MIT许可证发布。