NVlabs/FoundationPose

GitHub: NVlabs/FoundationPose

FoundationPose 是一个统一的 6D 物体位姿估计与跟踪基础模型,仅需 CAD 模型或少量参考图片即可对新物体实现免微调的高精度位姿估计。

Stars: 3409 | Forks: 507

# FoundationPose:统一的新型物体 6D 位姿估计与跟踪 [[论文]](https://arxiv.org/abs/2312.08344) [[网站]](https://nvlabs.github.io/FoundationPose/) 这是我们即将发表于 CVPR 2024 (Highlight) 的论文的官方实现。 我们提出了 FoundationPose,这是一个用于 6D 物体位姿估计和跟踪的统一基础模型,同时支持基于模型(model-based)和无模型(model-free)的设定。只要提供其 CAD 模型,或者捕获少量参考图像,我们的方法就可以在测试时立即应用于新型物体,而无需进行微调。我们通过神经隐式表示(neural implicit representation)弥合了这两种设定之间的差距,从而实现了有效的新型视图合成,并使下游的位姿估计模块在相同的统一框架下保持不变。通过在大规模合成数据上训练,并在大语言模型(LLM)、基于 Transformer 的新型架构以及对比学习公式的辅助下,实现了强大的泛化能力。在涉及具有挑战性的场景和物体的多个公开数据集上进行的广泛评估表明,我们的统一方法大大优于针对每个任务专门设计的现有方法。此外,尽管减少了假设条件,它的效果甚至可以与实例级方法(instance-level methods)相媲美。 **🤖 对于 ROS 版本,请查看 [Isaac ROS 位姿估计](https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_pose_estimation),它拥有 TRT 快速推理和 C++ 加速。** \ **🥇 在全球 [BOP 排行榜](https://bop.felk.cvut.cz/leaderboards/pose-estimation-unseen-bop23/core-datasets/) 上排名第一(截至 2024/03),针对基于模型的新型物体位姿估计。** ## 演示 机器人应用: 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 # Bibtex ``` @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. 从[这里](https://drive.google.com/drive/folders/1DFezOAD0oD1BblsXVxqDsl8fj0qzB82i?usp=sharing)下载所有网络权重,并将它们放在 `weights/` 文件夹下。对于 refiner,您需要 `2023-10-28-18-33-37`。对于 scorer,您需要 `2024-01-11-20-02-45`。 2. [下载演示数据](https://drive.google.com/drive/folders/1pRyFmxYXmAnpku7nGRioZaKrVJtIsroP?usp=sharing)并将它们解压到 `demo_data/` 文件夹下。 3. [可选] 下载我们的大规模训练数据:["FoundationPose Dataset"](https://drive.google.com/drive/folders/1s4pB6p4ApfWMiMjmTXOFco8dHbNXikp-?usp=sharing) 4. [可选] 在[此处](https://drive.google.com/drive/folders/1PXXCOJqHXwQTbwPwPbGDN9_vLVe0XpFS?usp=sharing)下载我们预处理过的参考视图,以便运行无模型的少样本(few-shot)版本。 # 环境搭建方式 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 ``` 对于 4090 等较新的 GPU,请参考[这里](https://github.com/NVlabs/FoundationPose/issues/27)。 简而言之,请执行以下操作: ``` 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”](https://pytorch.org/get-started/locally/) 页面列出了正确的 `--index-url`(例如 `cu124` 适用于大多数当前的 NVIDIA 驱动程序)。示例: ``` python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 ``` 3. **安装 PyTorch3D 和 NVDiffRast**(从源代码编译;需要 CUDA toolkit 以获取 `nvcc`。将 `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 安装文档](https://kaolin.readthedocs.io/en/latest/notes/installation.html))。 # 运行基于模型的演示 路径已通过 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 ``` 要运行无模型的少样本(few-shot)版本,您首先需要训练 Neural Object Field。`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 ``` # 故障排除 - 对于 4090 等较新的 GPU,请参考[这里](https://github.com/NVlabs/FoundationPose/issues/27)。 - 有关在 Windows 上的设置,请参考[这里](https://github.com/NVlabs/FoundationPose/issues/148)。 - 如果您得到不合理的结构,请查看[这里](https://github.com/NVlabs/FoundationPose/issues/44#issuecomment-2048141043)和[这里](https://github.com/030422Lee/FoundationPose_manual) # 训练数据下载 我们的训练数据包含使用来自 GSO 和 Objaverse 的 3D 资产的场景,具有高度逼真的渲染质量和大量的域随机化(domain randomization)。每个数据点都包含 **RGB、深度图、物体位姿、相机位姿、实例分割、2D 边界框**。[[Google Drive]](https://drive.google.com/drive/folders/1s4pB6p4ApfWMiMjmTXOFco8dHbNXikp-?usp=sharing)。 - 解析包含外参(extrinsics)和内参(intrinsics)在内的相机参数 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 存在法律限制,我们无法发布基于扩散的纹理增强(diffusion-based texture augmented)数据,也无法发布使用这些数据的预训练权重。因此,我们发布了未在基于扩散的增强数据上进行过训练的版本。预计会有轻微的性能下降。 # 许可证 代码和数据在 NVIDIA Source Code License 下发布。版权所有 © 2024 NVIDIA Corporation。保留所有权利。 # 联系方式 如有问题,请联系 [Bowen Wen](https://wenbowen123.github.io/)。
标签:6D姿态估计, Vectored Exception Handling, 凭据扫描, 增强现实, 机器人, 深度学习, 目标跟踪, 计算机视觉, 请求拦截, 逆向工具