YOLO26N 姿态估计 ONNX 导出与模型简化

YOLO26N 姿态估计 ONNX 导出与模型简化

1. ONNX 导出

#!/usr/bin/env python3"""export_pose_onnx.py"""fromultralyticsimportYOLO model=YOLO("runs/pose/yolo26n_pose_custom/weights/best.pt")# 导出 ONNXmodel.export(format="onnx",imgsz=640,opset=11,# 兼容性最好simplify=True,# 简化模型dynamic=False,# 边缘设备推荐静态batch=1,)# 输出：yolo26n-pose.onnx (~6.4MB)

2. ONNX 模型验证

#!/usr/bin/env python3"""verify_onnx.py"""importonnximportonnxruntimeasortimportnumpyasnp# 加载模型model=onnx.load("yolo26n-pose.onnx")onnx.checker.check_model(model)# 检查输入输出print("输入:")forinpinmodel.graph.input:shape=[d.dim_valuefordininp.type.tensor_type.shape.dim]print(f"{inp.name}:{shape}")print("输出:")foroutinmodel.graph.output:shape=[d.dim_valuefordinout.type.tensor_type.shape.dim]print(f"{out.name}:{shape}")# 推理测试session=ort.InferenceSession("yolo26n-pose.onnx")dummy=np.random.randn(1,3,640,640).astype(np.float32)outputs=session.run(None,{"images":dummy})print(f"输出 shape:{outputs[0].shape}")# [1, 56, 8400]

3. 模型简化（onnxsim）

# 安装 onnxsimpipinstallonnxsim# 简化模型onnxsim yolo26n-pose.onnx yolo26n-pose-sim.onnx# 验证简化后的模型python-c" import onnx model = onnx.load('yolo26n-pose-sim.onnx') print(f'节点数: {len(model.graph.node)}') print(f'大小: {os.path.getsize(\"yolo26n-pose-sim.onnx\") / 1024 / 1024:.1f} MB') "

4. Netron 可视化

# 使用 Netron 可视化模型结构# 访问 https://netron.app# 或本地安装pip install netron python-c"import netron; netron.start('yolo26n-pose.onnx')"# 检查要点：# 1. 输入 shape: [1, 3, 640, 640]# 2. 输出 shape: [1, 56, 8400]# 3. 无动态维度# 4. 算子兼容目标平台

5. ONNX 推理

#!/usr/bin/env python3"""onnx_pose_inference.py"""importonnxruntimeasortimportnumpyasnpimportcv2classONNXPoseDetector:def__init__(self,onnx_path,conf_thresh=0.3):self.session=ort.InferenceSession(onnx_path)self.conf_thresh=conf_thresh self.input_name=self.session.get_inputs()[0].namedefpreprocess(self,image):img=cv2.resize(image,(640,640))img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)img=img.astype(np.float32)/255.0img=img.transpose(2,0,1)returnnp.expand_dims(img,axis=0)defdetect(self,image):blob=self.preprocess(image)outputs=self.session.run(None,{self.input_name:blob})# 输出: [1, 56, 8400]predictions=outputs[0][0].T# [8400, 56]# 分解boxes=predictions[:,:4]# bboxscores=predictions[:,4]# confkpts=predictions[:,6:]# 关键点 (51 维)# 过滤mask=scores>self.conf_thresh boxes=boxes[mask]scores=scores[mask]kpts=kpts[mask]# 解析关键点results=[]foriinrange(len(boxes)):keypoints=kpts[i].reshape(17,3)results.append({'bbox':boxes[i],'score':scores[i],'keypoints':keypoints,})returnresultsif__name__=="__main__":detector=ONNXPoseDetector("yolo26n-pose.onnx")image=cv2.imread("test.jpg")results=detector.detect(image)print(f"检测到{len(results)}个人")

6. 输出解析

# YOLO26N-Pose 输出格式# tensor shape: [1, 56, 8400]# 56 = 4(bbox) + 1(conf) + 1(cls) + 17*3(kpts)defparse_pose_output(output,conf_thresh=0.3):"""解析姿态估计输出"""predictions=output[0].T# [8400, 56]# 边界框boxes=predictions[:,:4]# cx, cy, w, h# 置信度conf=predictions[:,4]# 类别（person=0）cls=predictions[:,5]# 关键点: [8400, 51] → [8400, 17, 3]kpts_raw=predictions[:,6:]kpts=kpts_raw.reshape(-1,17,3)# 过滤mask=conf>conf_threshreturnboxes[mask],conf[mask],kpts[mask]

总结