RealmanRobot/Foundation_Pose

FoundationPose Weights

这是 FoundationPose SDK 使用的模型权重仓库。

模型根据 RGB 图像、深度图、物体 mask、CAD 三维模型和相机内参，估计物体相对于相机坐标系的 6D 位姿。

本仓库只存放模型权重，不包含完整的推理代码。推理代码和环境配置请参见配套 SDK 代码仓库。

配套 SDK 代码

请先从配套 GitHub 代码仓库下载 FoundationPose SDK，并按照该仓库的说明完成环境配置：

FoundationPose SDK GitHub 仓库

本 Hugging Face 仓库用于存放模型权重，GitHub 仓库用于存放推理代码和运行环境。两者需要配套使用。

文件说明

本仓库根目录包含两个权重文件：

predict_ckpt.pth
refine_ckpt.pth

• predict_ckpt.pth：预测阶段权重。
• refine_ckpt.pth：位姿 refinement 阶段权重。

两个权重需要同时使用，不能只加载其中一个。

权重放置方式

从本仓库下载 predict_ckpt.pth 和 refine_ckpt.pth 后，按照 GitHub SDK 代码中的目录约定放置：

<FoundationPose SDK代码目录>/
└── tests/
    └── weights/
        ├── predict_ckpt/
        │   └── predict_ckpt.pth
        └── refine_ckpt/
            └── refine_ckpt.pth

如果不使用上述目录，也可以在调用 SDK 时直接传入两个权重文件的本地路径：

predict_ckpt_path = "/path/to/predict_ckpt.pth"
refine_ckpt_path = "/path/to/refine_ckpt.pth"

请以配套 GitHub SDK 代码的实际接口和目录约定为准。

接口示例

完成 GitHub SDK 的安装，并按照上面的目录结构放置两个权重后，可以使用以下代码进行位姿估计。示例默认使用 tests/demo_data/test_img 中的 RGB 图像、深度图、mask 和相机内参，以及一个 CAD 网格模型。

import copy
import json
import os.path
from types import SimpleNamespace

import cv2
import pyrealsense2 as rs

from FoundationPose.estimater11 import *
from FoundationPose.datareader import *
from FoundationPose.foundationpose_main import Detect_foundationpose


def load_resources(mesh_path, intrinsics, predict_ckpt_dir, refine_ckpt_dir):
    est, reader, bbox, debug, to_origin = (
        Detect_foundationpose.load_model(
            mesh_path,
            intrinsics,
            predict_ckpt_dir,
            refine_ckpt_dir,
        )
    )
    return est, reader, bbox, debug, to_origin


def main():
    image_path = "tests/demo_data/test_img"
    color_path = os.path.join(image_path, "rgb.png")
    depth_path = os.path.join(image_path, "depth.png")
    mask_path = os.path.join(image_path, "mask.png")

    mesh_path = "tests/demo_data/haoliyou/mesh/textured_mesh.obj"

    predict_ckpt_dir = "tests/weights/predict_ckpt/predict_ckpt.pth"
    refine_ckpt_dir = "tests/weights/refine_ckpt/refine_ckpt.pth"

    json_file = os.path.join(image_path, "intrinsics.json")
    with open(json_file, "r+") as fp:
        intrinsics = json.load(
            fp,
            object_hook=lambda value: SimpleNamespace(**value),
        )

    color_img = cv2.imread(color_path)
    depth_img = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

    # 加载 CAD 模型、相机内参和两个模型权重。
    est, reader, bbox, debug, to_origin = load_resources(
        mesh_path,
        intrinsics,
        predict_ckpt_dir,
        refine_ckpt_dir,
    )

    # 根据 RGB 图像、深度图和 mask 估计物体位姿。
    pose, color, to_origin = Detect_foundationpose.pose_est(
        color_img,
        depth_img,
        mask,
        reader,
        est,
        to_origin,
        bbox,
        show=True,
    )

    # 显示位姿可视化结果。
    color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)
    cv2.imshow("pose", color)
    cv2.waitKey(0)


if __name__ == "__main__":
    main()

其中：

• mesh_path 是待估计物体的 CAD 网格模型路径；
• predict_ckpt_dir 和 refine_ckpt_dir 是两个模型权重文件路径；
• color_img、depth_img 和 mask 是推理输入；
• intrinsics 是相机内参；
• pose 是模型输出的物体 6D 位姿；
• color 是包含位姿可视化结果的图像。

请根据实际数据路径、CAD 模型路径和相机内参文件修改示例中的路径。

使用要求

推理时还需要：

• RGB 图像；
• 深度图；
• 物体分割 mask；
• 对应物体的 CAD 三维模型；
• 相机内参；
• 与权重版本匹配的 FoundationPose SDK 代码和运行环境。

输出为物体的 6D 位姿。位姿矩阵的坐标系方向和单位请以配套 SDK 的接口说明为准。

注意事项

• 模型需要配套的 CAD 模型和 mask 才能进行推理。
• 实际效果会受到相机标定、深度噪声、物体遮挡、反光材质和 CAD 模型精度影响。
• 请在正式使用前，根据目标相机和物体进行独立测试。
• 未经验证，不建议将结果直接用于人身安全相关或高风险控制系统。

许可证信息

本项目遵循 MIT 许可证。