qing
ab60b3cafb
零 torch 依赖（hailo_platform + numpy + PIL），vstream 形状/量化参数从 HEF 动态推导，支持 uint8 量化边缘与动作归一化 stats；新增 6 个免硬件 单测（fake hailo_platform 注入）。调研报告记录 hailort 5.x 新 API 踩坑， 测速报告给出 CPU(ORT 5.9ms) vs NPU(resnet50 4.29ms) 基准与部署建议。
251 lines
10 KiB
Python
251 lines
10 KiB
Python
"""actpred HEF（Hailo）推理封装 —— AI HAT 2 / Hailo-10H 部署后端
|
||
|
||
接口对齐 predict.ActionPredictorInference：predict(images, action_history)。
|
||
|
||
HEF 由 actpred ONNX 经 Hailo DFC 编译生成（量化校准后）。本封装不硬编码
|
||
vstream 名称/形状/量化参数，全部从 HEF 的 input/output vstream info 动态推导：
|
||
|
||
图像输入: (1, F, 3, S, S) → NHWC uint8 (1, F, S, S, 3)
|
||
或 (1, S, S, F*3) → 帧沿通道拼接
|
||
动作输入: (1, H, D) float32，uint8 边缘时按 quant_info 量化
|
||
动作输出: (1, D)，uint8 边缘时按 quant_info 反量化
|
||
|
||
背景与踩坑记录见
|
||
docs/reports/actpred_hailo_hat2_inference_research_20260813.md：
|
||
- hailort 5.x 对 Hailo-10 仅支持新 API（create_infer_model），
|
||
旧版 VDevice.configure(hef) 报 HAILO_NOT_IMPLEMENTED；
|
||
- 绑定缓冲区 dtype/大小必须匹配 HEF 边缘格式（量化 uint8）；
|
||
- 图像输入按 Hailo 惯例为 uint8 原始像素（ImageNet 归一化由
|
||
DFC 编译时折叠进网络）。
|
||
"""
|
||
|
||
import argparse
|
||
import contextlib
|
||
import json
|
||
import os
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
# 与 predict.parse_action_history 相同逻辑；此处内联以保持
|
||
# 本模块零 torch 依赖（HAT 部署环境只装 hailo_platform + numpy + PIL）
|
||
|
||
|
||
def parse_action_history(s: str, action_dim: int) -> np.ndarray:
|
||
steps = s.split(":")
|
||
history = []
|
||
for step in steps:
|
||
values = [float(v) for v in step.split(",")]
|
||
if len(values) != action_dim:
|
||
raise ValueError(
|
||
f"Expected {action_dim} values per step, got {len(values)}")
|
||
history.append(values)
|
||
return np.array(history, dtype=np.float32)
|
||
|
||
|
||
class HailoPredictor:
|
||
"""HEF 动作预测推理器（Hailo NPU 后端）
|
||
|
||
用法：
|
||
with HailoPredictor("actpred.hef") as predictor:
|
||
action = predictor.predict(images, action_history)
|
||
"""
|
||
|
||
def __init__(self, hef_path: str, stats_path: str = None, vdevice=None):
|
||
from hailo_platform import HEF, VDevice
|
||
|
||
self._stack = contextlib.ExitStack()
|
||
if vdevice is None:
|
||
vdevice = self._stack.enter_context(VDevice())
|
||
self._vdevice = vdevice
|
||
self._infer_model = self._vdevice.create_infer_model(str(hef_path))
|
||
self._configured = self._stack.enter_context(
|
||
self._infer_model.configure())
|
||
try:
|
||
hef = HEF(str(hef_path))
|
||
self._in_infos = hef.get_input_vstream_infos()
|
||
self._out_infos = hef.get_output_vstream_infos()
|
||
if len(self._out_infos) != 1:
|
||
raise ValueError(
|
||
f"HEF 应只有一个输出 vstream，实际 {len(self._out_infos)}")
|
||
self._image_info = self._find_image_input()
|
||
self._history_info = self._find_history_input()
|
||
self._out_info = self._out_infos[0]
|
||
self._derive_dims()
|
||
|
||
self._action_mean = None
|
||
self._action_std = None
|
||
if stats_path is not None:
|
||
if not os.path.exists(stats_path):
|
||
raise FileNotFoundError(f"统计文件不存在: {stats_path}")
|
||
stats = np.load(stats_path)
|
||
if stats.shape != (2, self.action_dim):
|
||
raise ValueError(
|
||
f"action_stats.npy 形状应为 (2, {self.action_dim})，"
|
||
f"实际为 {stats.shape}: {stats_path}")
|
||
self._action_mean = stats[0].astype(np.float32)
|
||
self._action_std = stats[1].astype(np.float32)
|
||
except Exception:
|
||
self._stack.close() # 初始化失败也释放设备，避免 fd 泄漏
|
||
raise
|
||
|
||
# ---------------- 形状推导 ----------------
|
||
|
||
def _find_image_input(self):
|
||
cands = [i for i in self._in_infos if len(i.shape) in (4, 5)]
|
||
if len(cands) != 1:
|
||
raise ValueError(
|
||
"HEF 输入中无法唯一识别图像 vstream（期望一个 4/5 维输入），"
|
||
f"实际: {[(i.name, tuple(i.shape)) for i in self._in_infos]}")
|
||
return cands[0]
|
||
|
||
def _find_history_input(self):
|
||
cands = [i for i in self._in_infos
|
||
if len(i.shape) == 3 and i is not self._image_info]
|
||
if len(cands) != 1:
|
||
raise ValueError(
|
||
"HEF 输入中无法唯一识别动作历史 vstream（期望一个 3 维输入），"
|
||
f"实际: {[(i.name, tuple(i.shape)) for i in self._in_infos]}")
|
||
return cands[0]
|
||
|
||
def _derive_dims(self):
|
||
shape = tuple(self._image_info.shape)
|
||
if len(shape) == 5:
|
||
_, self.frame_len, c, self.image_size, _ = shape
|
||
if c != 3:
|
||
raise ValueError(f"图像输入通道维应为 3，实际形状 {shape}")
|
||
self._image_layout = "ncfhw"
|
||
else: # (1, S, S, F*3) NHWC 帧拼通道
|
||
_, self.image_size, s2, c = shape
|
||
if s2 != self.image_size or c % 3 != 0 or c == 3:
|
||
raise ValueError(
|
||
f"无法从 4 维图像输入推导帧数（期望 (1,S,S,F*3)），"
|
||
f"实际形状 {shape}")
|
||
self.frame_len = c // 3
|
||
self._image_layout = "nhwc_fused"
|
||
|
||
h_shape = tuple(self._history_info.shape)
|
||
self.history_len, self.action_dim = h_shape[1], h_shape[2]
|
||
|
||
out_shape = tuple(self._out_info.shape)
|
||
out_dim = out_shape[-1]
|
||
if out_dim != self.action_dim:
|
||
raise ValueError(
|
||
f"输出动作维度 {out_dim} 与动作历史维度 {self.action_dim} 不一致")
|
||
|
||
# ---------------- 生命周期 ----------------
|
||
|
||
def close(self):
|
||
self._stack.close()
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
self.close()
|
||
return False
|
||
|
||
# ---------------- 推理 ----------------
|
||
|
||
def predict(self, images, action_history: np.ndarray) -> np.ndarray:
|
||
"""images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧）
|
||
|
||
返回 (action_dim,) float32 动作向量。
|
||
"""
|
||
if not isinstance(images, (list, tuple)):
|
||
images = [images]
|
||
if len(images) != self.frame_len:
|
||
raise ValueError(
|
||
f"images 应为 {self.frame_len} 帧，实际为 {len(images)} 帧。"
|
||
"冷启动时可重复传入同一帧。")
|
||
action_history = np.asarray(action_history, dtype=np.float32)
|
||
expected = (self.history_len, self.action_dim)
|
||
if action_history.shape != expected:
|
||
raise ValueError(
|
||
f"action_history 形状应为 {expected}，实际为 {action_history.shape}")
|
||
|
||
if self._action_mean is not None:
|
||
action_history = (action_history - self._action_mean) / self._action_std
|
||
|
||
from hailo_platform import FormatType
|
||
bindings = self._configured.create_bindings()
|
||
|
||
frames = np.stack([self._preprocess(img) for img in images])
|
||
bindings.input(self._image_info.name).set_buffer(
|
||
self._image_buffer(frames))
|
||
|
||
h_shape = tuple(self._history_info.shape)
|
||
if self._history_info.format.type == FormatType.UINT8:
|
||
q = self._history_info.quant_info
|
||
buf = np.clip(np.round(action_history / q.qp_scale + q.qp_zp),
|
||
0, 255).astype(np.uint8).reshape(h_shape)
|
||
else:
|
||
buf = action_history.astype(np.float32).reshape(h_shape)
|
||
bindings.input(self._history_info.name).set_buffer(buf)
|
||
|
||
out_shape = tuple(self._out_info.shape)
|
||
if self._out_info.format.type == FormatType.UINT8:
|
||
out_buf = np.zeros(out_shape, dtype=np.uint8)
|
||
else:
|
||
out_buf = np.zeros(out_shape, dtype=np.float32)
|
||
bindings.output(self._out_info.name).set_buffer(out_buf)
|
||
|
||
self._configured.run([bindings], timeout=10000)
|
||
|
||
result = np.asarray(
|
||
bindings.output(self._out_info.name).get_buffer(),
|
||
dtype=np.float32)
|
||
if self._out_info.format.type == FormatType.UINT8:
|
||
q = self._out_info.quant_info
|
||
result = (result - q.qp_zp) * q.qp_scale
|
||
result = result.reshape(-1)
|
||
if self._action_mean is not None:
|
||
result = result * self._action_std + self._action_mean
|
||
return result
|
||
|
||
def _preprocess(self, image) -> np.ndarray:
|
||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||
img = img.convert("RGB").resize(
|
||
(self.image_size, self.image_size), Image.BILINEAR)
|
||
return np.asarray(img, dtype=np.uint8) # HWC uint8
|
||
|
||
def _image_buffer(self, frames: np.ndarray) -> np.ndarray:
|
||
"""(F, S, S, 3) → HEF 图像输入缓冲区（按编译后布局）"""
|
||
if self._image_layout == "ncfhw":
|
||
return frames.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,S,S)
|
||
# nhwc_fused: 帧沿通道拼接 (1, S, S, F*3)
|
||
return frames.reshape(1, self.image_size, self.image_size,
|
||
self.frame_len * 3)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="actpred HEF (Hailo) inference")
|
||
parser.add_argument("--hef", required=True, help="HEF 模型路径")
|
||
parser.add_argument("--images", required=True, nargs="+",
|
||
help="按时间顺序（旧→新）传入 frame_len 个图像路径")
|
||
parser.add_argument("--action-history", required=True,
|
||
help="冒号分隔各步，每步 action_dim 个值")
|
||
parser.add_argument("--stats", default=None,
|
||
help="action_stats.npy 路径（训练时启用归一化才需要）")
|
||
parser.add_argument("--output-format", choices=["json", "raw"],
|
||
default="json")
|
||
args = parser.parse_args()
|
||
|
||
with HailoPredictor(args.hef, stats_path=args.stats) as predictor:
|
||
if len(args.images) != predictor.frame_len:
|
||
raise SystemExit(
|
||
f"--images 需传入 {predictor.frame_len} 个路径（旧→新），"
|
||
f"实际为 {len(args.images)} 个")
|
||
images = [np.array(Image.open(p).convert("RGB")) for p in args.images]
|
||
history = parse_action_history(args.action_history, predictor.action_dim)
|
||
predicted = predictor.predict(images, history)
|
||
|
||
if args.output_format == "json":
|
||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||
else:
|
||
print(" ".join(f"{v:.4f}" for v in predicted))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|