使用 ONNX 加载头部姿态评估模型并集成到 LLM Agent

根据模型定义，输入要求为 [N, 3, 224, 224]，其中 N 是批次大小，3 是 RGB 通道，224x224 是图像尺寸。输出为 [N, 3]，分别对应 Yaw、Pitch、Roll 的角度值。

4. 推理实现

使用 onnxruntime 进行推理比直接使用 onnx 更高效。它支持多种执行提供者（Execution Provider），如 CPU、CUDA、CoreML 等。

4.1 初始化会话

import onnxruntime as ort

# 创建推理会话，默认使用 CPU，也可指定 CUDA 或 CoreML
session = ort.InferenceSession(
    "models/sixdrepnet360_Nx3x224x224.onnx", 
    providers=['CPUExecutionProvider']
)

# 获取输入输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

4.2 随机数据测试

在正式处理图片前，可以先用随机数据验证推理流程是否通畅。

import numpy as np

# 构造随机输入数据 (Batch=1, Channels=3, H=224, W=224)
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 执行推理
outputs = session.run([output_name], {input_name: input_data})
print(f"Inference Result: {outputs[0]}")

4.3 真实图像预处理与推理

实际场景中，需要读取图像并进行标准化预处理。6DRepNet 要求输入图像经过特定的归一化处理。

import cv2
import numpy as np

def preprocess_image(image_path):
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found: {image_path}")
    
    # 转换为 RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 调整大小至 224x224
    resized_image = cv2.resize(image, (224, 224))
    
    # 数据类型转换与归一化
    resized_image = resized_image.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    resized_image = (resized_image - mean) / std
    
    # 调整形状从 (H, W, C) 到 (C, H, W)
    resized_image = np.transpose(resized_image, (2, 0, 1))
    
    # 增加批次维度 (1, C, H, W)
    resized_image = np.expand_dims(resized_image, axis=0)
    
    return resized_image

# 执行推理
image_tensor = preprocess_image('images/test_face.jpg')
outputs = session.run([output_name], {input_name: image_tensor})
yaw, pitch, roll = outputs[0][0]
print(f"Head Pose - Yaw: {yaw:.2f}, Pitch: {pitch:.2f}, Roll: {roll:.2f}")

注意：在实际应用中，如果图像不是正面人脸，建议先进行人脸检测（如 SCRFD 模型）并裁剪出人脸区域，再进行上述预处理，以提高准确率。

5. 集成到 LLM Agent

为了让大语言模型能够调用此功能，我们可以使用 autogen 库将推理函数注册为工具。

5.1 定义工具函数

from typing import Annotated, TypedDict
import autogen

class HeadPose(TypedDict):
    Pitch: Annotated[float, "the pitch of head"]
    Yaw: Annotated[float, "the yaw of head"]
    Roll: Annotated[float, "the roll of head"]

def estimate_head_pose(image_path: Annotated[str, "the image path to be estimate head pose"]) -> Annotated[
    HeadPose, "A dictionary representing a head pose with pitch yaw roll."]:
    """Estimate the head pose of a given image using ONNX model."""
    try:
        input_name = session.get_inputs()[0].name
        output_name = session.get_outputs()[0].name
        
        # 复用预处理逻辑
        image = cv2.imread(image_path)
        if image is None:
            return {"Pitch": 0.0, "Yaw": 0.0, "Roll": 0.0}
        
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        resized_image = cv2.resize(image, (224, 224))
        resized_image = resized_image.astype(np.float32) / 255.0
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        resized_image = (resized_image - mean) / std
        resized_image = np.transpose(resized_image, (2, 0, 1))
        resized_image = np.expand_dims(resized_image, axis=0)
        resized_image = resized_image.astype(np.float32)
        
        outputs = session.run([output_name], {input_name: resized_image})
        result = outputs[0][0].tolist()
        
        return {
            "Pitch": float(result[1]), 
            "Yaw": float(result[0]), 
            "Roll": float(result[2])
        }
    except Exception as e:
        return {"Pitch": 0.0, "Yaw": 0.0, "Roll": 0.0}

5.2 注册工具并对话

# 假设 assistant 和 user_proxy 已初始化
autogen.register_function(
    estimate_head_pose, 
    caller=assistant, 
    executor=user_proxy,
    description="Estimate the head pose of a given image"
)

# 发起对话
chat_result = user_proxy.initiate_chat(
    assistant, 
    message="告诉我这个图像中人的头部姿态是什么？请提供图片路径。"
)

Agent 会自动识别意图，提取图片路径，调用函数，并将结果以自然语言反馈给用户。

6. 常见问题与最佳实践

1. 精度问题：在图像模糊或光照极差的情况下，姿态估算准确度可能下降。建议在预处理前增加图像增强步骤。
2. 执行提供者：若需加速推理，可尝试启用 GPU 支持（如 CUDAExecutionProvider），但需确保驱动和运行时版本匹配。
3. 参数解释：Yaw 代表左右转动，Pitch 代表上下点头，Roll 代表倾斜。Agent 应被提示不要捏造参数，若无法获取则询问用户。
4. 依赖管理：不同操作系统下 onnxruntime 的安装可能需要特定编译选项，建议使用预编译 wheel 包。

7. 总结

本文详细演示了如何将 6DRepNet 头部姿态评估模型转换为 ONNX 格式，并利用 ONNX Runtime 在 Python 环境中完成推理。通过 OpenCV 进行图像预处理，结合 AutoGen 将模型封装为 Agent 工具，实现了 LLM 对视觉信息的感知能力。这种方法降低了非算法工程师使用视觉模型的门槛，有助于丰富 LLM Agent 的工具栈，使其能够处理更复杂的现实世界任务。