Modelpublic

RealmanRobot/Foundation_Pose

sourceHugging Facemitupdated 6h agoView on Hugging Face

0likes

Model Card

FoundationPose Weights

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

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

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

配套 SDK 代码

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

FoundationPose SDK GitHub 仓库

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

文件说明

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

text

predict_ckpt.pth
refine_ckpt.pth

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

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

权重放置方式

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

text

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

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

python

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 网格模型。

python

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 许可证。