feat(pybot×actpred): R2 姿态反馈方案一——MPU6050 独立 IMU token 全链路
- pybot：recorder 逐帧落盘 imu.npy（T×8，roll,pitch,ax,ay,az,gx,gy,gz，与帧/动作同调用对齐）；run_loop 对声明 imu_dim 的 predictor 以 predict(..., imu=) 调用（鸭子类型，旧两参调用向后兼容） - actpred：新增 --imu-dim；model 序列 [img×F, act×H, (IMU), PRED]（imu_dim>0 时 F+H+2，IMU token 位于 PRED 前）；dataset/train 读 imu.npy（缺失按零向量兜底）；predict/ONNX/Hailo 三后端与 export/gen_calibration/dfc_compile/_npu_test 部署链适配 imu 输入 - 姿态按物理量程固定缩放归一化（imu_norm.py，训练/推理共用，零状态） - 测试：actpred 128 passed / pybot loop+recorder 28 passed
This commit is contained in:
qing committed
1 parent
3e773debc9
commit
a2ba797410
22 files changed
+1236
-133
No files matched your search
+33
-10
@@ -1,10 +1,11 @@
|
||||
"""actpred ONNX Runtime 推理封装 —— CPU 后端（Pi 5 实测 5.9ms/次，PyTorch 的 ~5 倍）
|
||||
|
||||
接口对齐 predict.ActionPredictorInference：predict(images, action_history)。
|
||||
接口对齐 predict.ActionPredictorInference：predict(images, action_history, imu=None)。
|
||||
|
||||
模型维度从 ONNX session 输入元数据动态推导（images: [1,F,3,H,W]，
|
||||
action_history: [1,H,D]），无需 checkpoint。
|
||||
图像预处理复现训练时的 dataset.py 流水线：Resize → /255 → ImageNet 归一化。
|
||||
action_history: [1,H,D]；R2 imu_dim>0 的模型另有 imu: [1,imu_dim]），
|
||||
无需 checkpoint。图像预处理复现训练时的 dataset.py 流水线：
|
||||
Resize → /255 → ImageNet 归一化。
|
||||
|
||||
背景见 docs/reports/actpred_inference_speed_benchmark_20260813.md。
|
||||
"""
|
||||
@@ -16,6 +17,8 @@ import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from imu_norm import normalize_imu
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
@@ -54,11 +57,16 @@ class OnnxPredictor:
|
||||
|
||||
inputs = {i.name: [d if isinstance(d, int) else -1
|
||||
for d in i.shape] for i in self._session.get_inputs()}
|
||||
img_shape, hist_shape = self._locate_inputs(inputs)
|
||||
img_shape, hist_shape, imu_shape = self._locate_inputs(inputs)
|
||||
_, self.frame_len, _, self.image_height, self.image_width = img_shape
|
||||
_, self.history_len, self.action_dim = hist_shape
|
||||
self._image_name = next(n for n, s in inputs.items() if len(s) == 5)
|
||||
self._history_name = next(n for n, s in inputs.items() if len(s) == 3)
|
||||
# R2 方案一：2 维输入 [1, imu_dim] 为姿态（无则 None）。
|
||||
# pybot 闭环鸭子类型约定：None 时按旧两参调用。
|
||||
self.imu_dim = imu_shape[-1] if imu_shape is not None else None
|
||||
self._imu_name = (next(n for n, s in inputs.items() if len(s) == 2)
|
||||
if imu_shape is not None else None)
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
@@ -75,17 +83,19 @@ class OnnxPredictor:
|
||||
|
||||
@staticmethod
|
||||
def _locate_inputs(inputs):
|
||||
"""按维度数识别图像(5维)/动作历史(3维)输入，返回各自形状"""
|
||||
"""按维度数识别图像(5维)/动作历史(3维)/姿态(2维,可选)输入，返回各自形状"""
|
||||
img = [s for s in inputs.values() if len(s) == 5]
|
||||
hist = [s for s in inputs.values() if len(s) == 3]
|
||||
if len(img) != 1 or len(hist) != 1:
|
||||
imu = [s for s in inputs.values() if len(s) == 2]
|
||||
if len(img) != 1 or len(hist) != 1 or len(imu) > 1:
|
||||
raise ValueError(
|
||||
"ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史，"
|
||||
"ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史"
|
||||
"（R2 可再加 1 个 2 维 imu），"
|
||||
f"实际: {inputs}")
|
||||
if img[0][2] != 3:
|
||||
raise ValueError(
|
||||
f"图像输入形状应为 [1,F,3,H,W]，实际为 {img[0]}")
|
||||
return img[0], hist[0]
|
||||
return img[0], hist[0], (imu[0] if imu else None)
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
@@ -101,8 +111,11 @@ class OnnxPredictor:
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
def predict(self, images, action_history: np.ndarray) -> np.ndarray:
|
||||
def predict(self, images, action_history: np.ndarray,
|
||||
imu: np.ndarray = None) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧）
|
||||
imu: 可选原始姿态向量（roll,pitch,ax,ay,az,gx,gy,gz，模型含 imu
|
||||
输入时生效，内部按物理量程归一化；None 等价零向量）
|
||||
|
||||
返回 (action_dim,) float32 动作向量。
|
||||
"""
|
||||
@@ -128,6 +141,9 @@ class OnnxPredictor:
|
||||
self._image_name: frames[np.newaxis], # (1,F,3,S,S)
|
||||
self._history_name: action_history[np.newaxis], # (1,H,D)
|
||||
}
|
||||
if self._imu_name is not None:
|
||||
feeds[self._imu_name] = normalize_imu(
|
||||
imu, self.imu_dim)[np.newaxis] # (1,imu_dim)
|
||||
result = self._session.run(None, feeds)[0].reshape(-1).astype(np.float32)
|
||||
if self._action_mean is not None:
|
||||
result = result * self._action_std + self._action_mean
|
||||
@@ -152,6 +168,9 @@ def main():
|
||||
help="冒号分隔各步，每步 action_dim 个值")
|
||||
parser.add_argument("--stats", default=None,
|
||||
help="action_stats.npy 路径（训练时启用归一化才需要）")
|
||||
parser.add_argument("--imu", type=str, default=None,
|
||||
help="可选姿态向量（逗号分隔，原始量纲：roll,pitch 度 / "
|
||||
"ax..gz int16 原始值）；模型含 imu 输入且缺省时按零向量")
|
||||
parser.add_argument("--output-format", choices=["json", "raw"],
|
||||
default="json")
|
||||
args = parser.parse_args()
|
||||
@@ -163,7 +182,11 @@ def main():
|
||||
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)
|
||||
imu = None
|
||||
if args.imu is not None:
|
||||
imu = np.array([float(v) for v in args.imu.split(",")],
|
||||
dtype=np.float32)
|
||||
predicted = predictor.predict(images, history, imu=imu)
|
||||
|
||||
if args.output_format == "json":
|
||||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||||
|
||||
Reference in new issue
Block a user