FoundationPose：统一的新物体6D姿态估计和跟踪

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

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

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

🤖 ROS版本请查看Isaac ROS姿态估计，它享有TRT快速推理和C++加速。

🥇 在全球BOP排行榜上排名第一（截至2024/03），用于基于模型的新物体姿态估计。 <img src="https://yellow-cdn.veclightyear.com/835a84d5/55a177d1-2c5f-4196-8635-1e52c1622cc6.jpg" width="80%">

演示

机器人应用：

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

AR应用：

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},
}

数据准备

从这里下载所有网络权重并将它们放在weights/文件夹下
下载演示数据并将它们解压到demo_data/文件夹下
[可选] 下载我们的大规模训练数据："FoundationPose数据集"](https://drive.google.com/drive/folders/1s4pB6p4ApfWMiMjmTXOFco8dHbNXikp-?usp=sharing)
[可选] 下载我们预处理的参考视图在这里以运行无模型少样本版本。

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

cd docker/
docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose  # 或者从头开始构建：docker build --network host -t foundationpose .
bash docker/run_container.sh

如果这是您第一次启动容器，您需要构建扩展。

bash build_all.sh

之后您可以无需重新构建即可执行进入容器。

docker exec -it foundationpose bash

环境设置选项2：conda（实验性）

设置conda环境

# 创建conda环境
conda create -n foundationpose python=3.9

# 激活conda环境
conda activate foundationpose

# 在conda环境下安装Eigen3 3.4.0
conda install conda-forge::eigen=3.4.0
export CMAKE_PREFIX_PATH="$CMAKE_PREFIX_PATH:/eigen/path/under/conda"

# 安装依赖
python -m pip install -r requirements.txt

# 安装NVDiffRast
python -m pip install --quiet --no-cache-dir git+https://github.com/NVlabs/nvdiffrast.git

# Kaolin（可选，如果运行无模型设置则需要）
python -m pip install --quiet --no-cache-dir kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-2.0.0_cu118.html

# PyTorch3D
python -m pip install --quiet --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu118_pyt200/download.html

# 构建扩展
CMAKE_PREFIX_PATH=$CONDA_PREFIX/lib/python3.9/site-packages/pybind11/share/cmake/pybind11 bash build_all_conda.sh

运行基于模型的演示

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

python run_demo.py

随意尝试其他物体（无需重新训练），比如电钻，只需更改argparse中的路径即可。 <img src="https://yellow-cdn.veclightyear.com/835a84d5/6da72d4c-4748-4a6b-aa82-bde5a935822a.jpg" width="50%">

在公开数据集上运行（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]。

解析相机参数（包括外参和内参）

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