FoundationPose: 新型物体的统一6D位姿估计与跟踪

[论文] [网站]

这是我们将在CVPR 2024（亮点论文）上发表的论文的官方实现。

贡献者：Bowen Wen、Wei Yang、Jan Kautz、Stan Birchfield

我们提出了FoundationPose，这是一个用于6D物体位姿估计与跟踪的统一基础模型，支持基于模型和无模型两种设置。我们的方法在测试时无需微调即可即时应用于新型物体，前提是提供该物体的CAD模型，或捕获少量参考图像。我们通过神经隐式表示弥合了这两种设置之间的差距，该表示支持有效的新视角合成，使下游位姿估计模块在同一统一框架下保持不变性。通过大规模合成训练，并借助大型语言模型（LLM）、新颖的基于Transformer的架构以及对比学习公式，实现了强大的泛化能力。在涉及具有挑战性场景和物体的多个公共数据集上进行的广泛评估表明，我们的统一方法显著优于针对各个任务专门设计的现有方法。此外，尽管减少了假设条件，它甚至达到了与实例级方法相当的结果。

🤖 关于ROS版本，请查看 Isaac ROS Pose Estimation，该版本具备TRT快速推理和C++加速功能。

🥇 在全球 BOP排行榜 上（截至2024年3月），基于模型的新型物体位姿估计任务排名第一。

演示

机器人应用：

https://github.com/NVlabs/FoundationPose/assets/23078192/aa341004-5a15-4293-b3da-000471fd74ed

增强现实应用：

https://github.com/NVlabs/FoundationPose/assets/23078192/80e96855-a73c-4bee-bcef-7cba92df55ca

YCB-Video数据集上的结果：

https://github.com/NVlabs/FoundationPose/assets/23078192/9b5bedde-755b-44ed-a973-45ec85a10bbe

参考文献格式

@InProceedings{foundationposewen2024,
author        = {Bowen Wen, Wei Yang, Jan Kautz, Stan Birchfield},
title         = {{FoundationPose}: Unified 6D Pose Estimation and Tracking of Novel Objects},
booktitle     = {CVPR},
year          = {2024},
}

如果您发现无模型设置有用，也请考虑引用：

@InProceedings{bundlesdfwen2023,
author        = {Bowen Wen and Jonathan Tremblay and Valts Blukis and Stephen Tyree and Thomas M\"{u}ller and Alex Evans and Dieter Fox and Jan Kautz and Stan Birchfield},
title         = {{BundleSDF}: {N}eural 6-{DoF} Tracking and {3D} Reconstruction of Unknown Objects},
booktitle     = {CVPR},
year          = {2023},
}

数据准备

1. 从此处下载所有网络权重，并将其放置在 weights/ 文件夹下。对于优化器（refiner），您需要 2023-10-28-18-33-37。对于评分器（scorer），您需要 2024-01-11-20-02-45。

2. 下载演示数据并将其解压到 demo_data/ 文件夹下。

3. [可选] 下载我们的大规模训练数据："FoundationPose Dataset"

4. [可选] 从此处下载我们预处理的参考视图，以运行无模型少样本版本。

环境设置选项 1：Docker（推荐）

  cd docker/
  docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose  # Or to build from scratch: docker build --network host -t foundationpose .
  bash docker/run_container.sh

如果是首次启动容器，您需要构建扩展。请在 Docker 容器内部运行以下命令。

bash build_all.sh

之后，您可以直接进入容器执行操作，无需重新构建。

docker exec -it foundationpose bash

对于较新型号的 GPU（如 4090），请参考 此链接。 简而言之，请执行以下操作：

docker pull shingarey/foundationpose_custom_cuda121:latest

然后修改 bash 脚本，使其使用此镜像而非 foundationpose:latest。

环境设置选项 2：conda（本地）

1. 创建环境（C++ 构建依赖项 + Python；均来自 conda-forge）：

conda env create -f environment.yml
conda activate foundationpose

2. 安装 PyTorch 时选择与您机器匹配的 CUDA 版本。PyTorch “Get Started” 页面列出了正确的 --index-url（例如，cu124 适用于大多数当前的 NVIDIA 驱动程序）。示例：

python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

3. 安装 PyTorch3D 和 NVDiffRast（从源代码编译；需要用于 nvcc 的 CUDA 工具包。将 CUDA_HOME 指向您的安装路径，例如 /usr/local/cuda-12.8 或 /usr/local/cuda）：

export CUDA_HOME=/usr/local/cuda   # or e.g. /usr/local/cuda-12.8
export PATH="$CUDA_HOME/bin:$PATH"
python -m pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git"
python -m pip install --no-build-isolation "git+https://github.com/NVlabs/nvdiffrast.git"

--no-build-isolation 是必需的，这样构建过程才能识别你已安装的 torch。

4. 安装剩余的 Python 依赖项并构建 mycpp 扩展（BundleSDF 的 mycuda 步骤为可选项；除非你需要无模型/NeRF 路径，否则可跳过）：

python -m pip install -r requirements.txt
bash build_all_conda.sh

5. 可选 — Kaolin（仅适用于无模型设置；版本必须与您的 PyTorch/CUDA 匹配—请参阅 Kaolin 安装文档）。

运行基于模型的演示

路径已在 argparse 中默认设置。如果需要更改场景，可以相应地传递参数。通过运行演示数据，您应该能够看到机器人操作芥末瓶。在第一帧进行位姿估计，然后自动切换到跟踪模式以处理视频的其余部分。生成的可视化结果将保存到 argparse 中指定的 debug_dir。（注意：首次运行可能会因在线编译而较慢）

python run_demo.py

欢迎尝试其他物体（无需重新训练），例如电钻，只需修改 argparse 中的路径即可。

在公开数据集（LINEMOD、YCB-Video）上运行

为此，您首先需要下载 LINEMOD 数据集和 YCB-Video 数据集。

要分别在这两个数据集上运行基于模型的版本，请根据您的下载位置设置路径。结果将保存到 debug 文件夹中。

python run_linemod.py --linemod_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD --use_reconstructed_mesh 0

python run_ycb_video.py --ycbv_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video --use_reconstructed_mesh 0

要运行无模型少样本版本，您首先需要训练神经对象场。ref_view_dir 基于您在上述“数据准备”部分下载数据的位置。将 dataset 标志设置为您感兴趣的数据集。

python bundlesdf/run_nerf.py --ref_view_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16 --dataset ycbv

然后运行与基于模型版本类似的命令，并稍作修改。这里我们以 YCB-Video 为例：

python run_ycb_video.py --ycbv_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video --use_reconstructed_mesh 1 --ref_view_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16

故障排除

• 对于较新型号的 GPU（如 4090），请参考 此链接。

• 关于在 Windows 系统上的设置，请参考 此链接。

• 如果您得到不合理的结果，请查看 此链接 和 此链接

训练数据下载

我们的训练数据包含使用来自 GSO 和 Objaverse 的 3D 资产构建的场景，这些场景通过高质量的照片级真实感渲染和大规模域随机化生成。每个数据点包括 RGB、深度、物体姿态、相机姿态、实例分割、2D 边界框。[Google Drive]。

• 解析包含外参和内参的相机参数

  glcam_in_cvcam = np.array([[1,0,0,0],
                          [0,-1,0,0],
                          [0,0,-1,0],
                          [0,0,0,1]]).astype(float)
  W, H = camera_params["renderProductResolution"]
  with open(f'{base_dir}/camera_params/camera_params_000000.json','r') as ff:
    camera_params = json.load(ff)
  world_in_glcam = np.array(camera_params['cameraViewTransform']).reshape(4,4).T
  cam_in_world = np.linalg.inv(world_in_glcam)@glcam_in_cvcam
  world_in_cam = np.linalg.inv(cam_in_world)
  focal_length = camera_params["cameraFocalLength"]
  horiz_aperture = camera_params["cameraAperture"][0]
  vert_aperture = H / W * horiz_aperture
  focal_y = H * focal_length / vert_aperture
  focal_x = W * focal_length / horiz_aperture
  center_y = H * 0.5
  center_x = W * 0.5
  
  fx, fy, cx, cy = focal_x, focal_y, center_x, center_y
  K = np.eye(3)
  K[0,0] = fx
  K[1,1] = fy
  K[0,2] = cx
  K[1,2] = cy
  

注意事项

由于在 LAION 数据集上训练的 Stable-Diffusion 存在法律限制，我们无法发布基于扩散模型的纹理增强数据，也无法发布使用该数据训练的预训练权重。因此，我们发布的版本未在扩散增强数据上进行训练，预计性能会有轻微下降。

致谢

感谢 Jeff Smith 在代码发布过程中提供的帮助；感谢 NVIDIA Isaac Sim 和 Omniverse 团队在合成数据生成方面给予的支持；感谢 Tianshi Cao 带来的富有价值的讨论。最后，我们也十分感谢 CVPR 评审专家和领域主席提出的积极反馈与建设性建议。

许可证

代码和数据基于 NVIDIA 源代码许可证发布。版权所有 © 2024，NVIDIA Corporation。保留所有权利。

联系方式

如有疑问，请联系 Bowen Wen。