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

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

随着大语言模型（LLM）能力的提升，将视觉模型作为工具嵌入到 LLM Agent 应用中已成为趋势。通过给 LLM 插上视觉模型的翅膀，可以让应用具备理解图像、分析场景的能力。本文将介绍如何使用 ONNX 格式加载头部姿态评估模型（6DRepNet），并通过 Python 代码实现推理，最后将其封装为工具函数集成到 AutoGen Agent 中。

1. 模型背景与原理

头部姿态估计是计算机视觉中的经典任务，通常输出 Yaw（偏航角）、Pitch（俯仰角）和 Roll（翻滚角）。传统的欧拉角表示存在万向节死锁问题，而论文《6DRepNet: A Simple yet Effective Approach for Head Pose Estimation》提出了一种基于 6D 旋转矩阵表示的方法，能够更稳健地回归完整的旋转外观。

该模型的核心优势包括：

• 无约束估计：支持全角度范围内的姿态预测。
• 连续表示：使用压缩的 6D 形式高效回归，避免离散分类误差。
• 几何感知损失：基于测地线距离的损失函数，符合 SO(3) 流形几何特性。

虽然 PyTorch 原生模型便于训练，但在部署阶段，ONNX（Open Neural Network Exchange）提供了更好的跨框架兼容性和性能优化，支持在云端、移动设备和嵌入式设备中运行。

2. 环境准备

首先，确保已安装必要的 Python 包。我们需要 onnx 用于模型解析，onnxruntime 用于推理，以及 opencv-python 和 numpy 用于图像处理。

pip install onnx onnxruntime opencv-python numpy

3. 加载与检查模型

ONNX 模型是一个黑盒，我们需要了解其输入和输出结构。可以使用 Netron 在线工具查看，也可以通过代码直接读取。

3.1 使用代码查看模型结构

import onnx

# 加载模型
model = onnx.load("models/sixdrepnet360_Nx3x224x224.onnx")

# 打印输入信息
print("Input:")
for inp in model.graph.input:
    print(f"  Name: {inp.name}, Type: {inp.type.tensor_type.elem_type}")
    if inp.type.tensor_type.HasField('shape'):
        dims = [d.dim_value if d.HasField('dim_value') else d.dim_param 
                for d in inp.type.tensor_type.shape.dim]
        print(f"  Shape: {dims}")

# 打印输出信息
print("Output:")
for out in model.graph.output:
    print(f"  Name: {out.name}, Type: {out.type.tensor_type.elem_type}")
    if out.type.tensor_type.HasField('shape'):
        dims = [d.dim_value if d.HasField('dim_value') else d.dim_param 
                for d in out.type.tensor_type.shape.dim]
        print(f"  Shape: {dims}")

根据模型定义，输入要求为 [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 的工具栈，使其能够处理更复杂的现实世界任务。