RK3588部署vitpose模型（最全方案）

目录前言一、GitHub下载源码二、去mmpose官网下载模型三、搭建环境四、pth转onnx五、onnx转rknn六、部署rknn模型模型效果总结前言在对比了vitpose和yolov11-pose两个模型发现对手的识别精度vitpose还是更高一些所以分享下在RK部署vitpose的经验一、GitHub下载源码去官网下载vitpose源码github.com/ViTAE-Transformer/ViTPose#二、去mmpose官网下载模型Body 2D Keypoint — MMPose 1.3.2 documentation三、搭建环境先是在ubuntu上搭建环境建议用conda虚拟环境这样不会影响到主环境主要依赖的库有以下这些mmcv、torchvision、matplotlib、timm、json_tricks、munkres搭建好环境就可以开始转换onnx了四、pth转onnx进入这个目录根据你们自己下载源码的目录修改成自己的目录cd /home/ViTPose-main运行转换代码from mmpose.apis.inference import init_pose_model from mmcv.runner import load_checkpoint import torch import types # -------------------------- 1. 路径配置 -------------------------- config_file /ViTPose-main/configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/ViTPose_small_coco_256x192.py checkpoint_file /ViTPose-main/td-hm_ViTPose-small_8xb64-210e_coco-256x192-62d7a712_20230314.pth onnx_save_path pose_model.onnx # -------------------------- 2. 加载模型 -------------------------- # 初始化空模型CPU模式 model init_pose_model(config_file, None, devicecpu) # 手动加载权重忽略参数名不匹配 load_checkpoint(model, checkpoint_file, map_locationcpu, strictFalse) # 强制设为推理模式关闭BatchNorm/Dropout model.eval() print(✅ 模型加载成功CPU模式) # -------------------------- 3. 重写forward方法适配所有ViTPose版本 -------------------------- def vitpose_forward_onnx(self, img): 适配MMPose框架的forward_dummy方法专门用于ONNX导出 forward_dummy是MMPose内置的无训练参数依赖的前向方法 # 调用框架内置的forward_dummy直接输出热力图 output self.forward_dummy(img) # 确保输出是单张量避免list/tuple导致ONNX导出异常 return output[0] if isinstance(output, (list, tuple)) else output # 替换模型的forward方法 model.forward types.MethodType(vitpose_forward_onnx, model) # -------------------------- 4. 先测试前向推理确保无错 -------------------------- try: dummy_input torch.randn(1, 3, 256, 192) with torch.no_grad(): test_output model(dummy_input) print(f✅ 前向推理测试成功输出形状{test_output.shape}) except Exception as e: print(f❌ 前向推理测试失败{e}) exit(1) # -------------------------- 5. 导出ONNX -------------------------- with torch.no_grad(): torch.onnx.export( model, dummy_input, # 复用测试的输入张量 onnx_save_path, input_names[input], # 输入节点名 output_names[heatmap], # 输出节点名关键点热力图 dynamic_axesNone, # 静态尺寸RK3588部署推荐 opset_version11, # 兼容RKNN的稳定版本 verboseFalse, # 关闭冗余日志 do_constant_foldingTrue, # 折叠常量节点简化模型 export_paramsTrue # 导出模型权重必须 ) print(f✅ ONNX模型导出完成保存路径{onnx_save_path}) # -------------------------- 6. 验证ONNX模型有效性 -------------------------- try: import onnx # 加载并检查模型语法 onnx_model onnx.load(onnx_save_path) onnx.checker.check_model(onnx_model) # 打印模型信息 input_shape onnx_model.graph.input[0].type.tensor_type.shape.dim output_shape onnx_model.graph.output[0].type.tensor_type.shape.dim print(f✅ ONNX模型验证通过) print(f 输入尺寸{[d.dim_value for d in input_shape]}) print(f 输出尺寸{[d.dim_value for d in output_shape]}) except ImportError: print(⚠️ 未安装onnx库跳过验证执行 pip install onnx 即可安装) except Exception as e: print(f❌ ONNX模型验证失败{str(e)[:200]})显示这样就代表成功转换了接下来就可以开始转换rknn了五、onnx转rknn首先要下载rknn_model_zoo-main和rknn-toolkit2-master这两个工具其中rknn_model_zoo-main为转换工具rknn-toolkit2-master为转换依赖的环境大家可以去瑞芯微的gitHub官网下载最新版本我在这贴我下载的版本路径github链接https://gitcode.com/gh_mirrors/rk/rknn-toolkit2/tree/v2.3.0?utm_sourcecsdn_github_acceleratorisLogin1--克隆V2.3.0标签github链接https://gitcode.com/gh_mirrors/rk/rknn_model_zoo/tree/v2.3.0?utm_sourcecsdn_github_acceleratorisLogin1--克隆V2.3.0标签具体安装步骤我在这篇就不写了大家可以参考我的另外一篇文章进行安装RK3588部署yolov11结合OpenCV进行实时推理保姆级建议收藏_rk3588 yolov11-CSDN博客模型转换代码from rknn.api import RKNN rknn RKNN(verboseTrue, verbose_file./conversion_log.txt) # 务必开启日志 # 配置转换参数这是精度和性能的调节阀 rknn.config( mean_values[[0, 0, 0]], # 预处理减去的均值 std_values[[255, 255, 255]], # 预处理除以的标准差 quantized_dtypeasymmetric_quantized-8, # 量化数据类型 quantized_algorithmnormal, # 量化算法 target_platformrk3588, # 目标平台 float_dtypefloat16, # NPU默认浮点类型 optimization_level3, # 优化等级越高优化越激进 ) # 加载ONNX模型 rknn.load_onnx(model./pose_model.onnx) # 构建RKNN模型 #ret rknn.build(do_quantizationTrue, dataset./dataset.txt) # 是否量化 ret rknn.build(do_quantizationFalse) if ret ! 0: print(Build failed!) exit(ret) # 导出模型 ret rknn.export_rknn(export_path./pose_model.rknn)执行命令/bin/python3 convert.py pose_model.onnx六、部署rknn模型注意照片路径要换成自己的照片路径然后执行以下的命令转换/bin/python3 infer_person.py这是源码import cv2 import numpy as np from rknnlite.api import RKNNLite import os import time # 新增导入时间模块 # 核心配置 INPUT_HEIGHT 256 INPUT_WIDTH 192 HEATMAP_HEIGHT INPUT_HEIGHT // 4 # 64 HEATMAP_WIDTH INPUT_WIDTH // 4 # 48 CONF_THRESH 0.1 # 临时降低阈值观察效果 # 初始化 RKNN-Lite rknn_lite RKNNLite() model_path pose_model.rknn if not os.path.exists(model_path): print(f❌ RKNN模型文件不存在{model_path}) exit(1) ret rknn_lite.load_rknn(model_path) if ret ! 0: print(❌ 加载RKNN模型失败!) exit(ret) ret rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_0_1_2) if ret ! 0: print(❌ 初始化运行环境失败!) exit(ret) print(✅ RKNN模型加载初始化成功) # 图像预处理你的原始函数保持不变 def preprocess_image(img_path): # 1. 检查图片路径 if not os.path.exists(img_path): raise FileNotFoundError(f❌ 图片不存在{img_path}) # 2. 读取图片兼容中文路径 img cv2.imdecode(np.fromfile(img_path, dtypenp.uint8), cv2.IMREAD_COLOR) if img is None: raise ValueError(f❌ 读取图片失败{img_path}文件损坏/格式不支持) orig_h, orig_w img.shape[:2] # 3. 等比例缩放 scale min(INPUT_WIDTH/orig_w, INPUT_HEIGHT/orig_h) new_w int(orig_w * scale) new_h int(orig_h * scale) img_resized cv2.resize(img, (new_w, new_h)) # 4. 填充到256×192 img_padded np.zeros((INPUT_HEIGHT, INPUT_WIDTH, 3), dtypenp.uint8) pad_x (INPUT_WIDTH - new_w) // 2 pad_y (INPUT_HEIGHT - new_h) // 2 img_padded[pad_y:pad_ynew_h, pad_x:pad_xnew_w, :] img_resized # 5. 预处理归一化 img_rgb cv2.cvtColor(img_padded, cv2.COLOR_BGR2RGB) img_float img_rgb.astype(np.float32) mean np.array([123.675, 116.28, 103.53], dtypenp.float32) std np.array([58.395, 57.12, 57.375], dtypenp.float32) img_norm (img_float - mean) / std img_trans img_norm.transpose(2, 0, 1) img_input np.expand_dims(img_trans, axis0) return img_input, orig_w, orig_h, pad_x, pad_y, scale # 推理 关键点解析修改版新增时间统计 def infer_keypoints(img_path): # 预处理预处理时间单独统计可选 try: input_img, orig_w, orig_h, pad_x, pad_y, scale preprocess_image(img_path) print(f✅ 图片预处理成功原图尺寸({orig_w}×{orig_h})缩放比例{scale:.2f}) except Exception as e: print(f❌ 预处理失败{e}) return [], 0.0 # 返回空关键点 0耗时 # ---------- 核心计时图片输入模型 → 2D坐标输出 开始 ---------- start_time time.perf_counter() # 高精度时间戳推荐 # 推理模型输入 outputs rknn_lite.inference(inputs[input_img]) heatmap outputs[0] # shape(1,17,64,48) print(f✅ 推理成功热力图形状{heatmap.shape}) # ---------- 调试打印热力图统计信息 ---------- print(f 热力图统计min{np.min(heatmap):.6f}, max{np.max(heatmap):.6f}, mean{np.mean(heatmap):.6f}) # ---------- 如果热力图数值过小如 0.01尝试使用 sigmoid 激活 ---------- if np.max(heatmap) 0.1: print(⚠️ 热力图最大值小于0.1自动应用 sigmoid 激活) heatmap 1.0 / (1.0 np.exp(-heatmap)) # sigmoid print(f 应用 sigmoid 后min{np.min(heatmap):.6f}, max{np.max(heatmap):.6f}) # 解析关键点2D坐标输出 keypoints [] for i in range(17): joint_heatmap heatmap[0, i, :, :] max_val np.max(joint_heatmap) if max_val CONF_THRESH: keypoints.append((0, 0, 0.0)) continue # 找到热力图中最大值的坐标 hm_y, hm_x np.unravel_index(np.argmax(joint_heatmap), (HEATMAP_HEIGHT, HEATMAP_WIDTH)) # 映射到输入图坐标 input_x hm_x * 4 input_y hm_y * 4 # 去除填充偏移并映射回原图坐标 input_x - pad_x input_y - pad_y orig_x int(round(input_x / scale)) orig_y int(round(input_y / scale)) # 边界保护 orig_x max(0, min(orig_x, orig_w - 1)) orig_y max(0, min(orig_y, orig_h - 1)) keypoints.append((orig_x, orig_y, float(max_val))) # ---------- 核心计时图片输入模型 → 2D坐标输出 结束 ---------- end_time time.perf_counter() total_time end_time - start_time # 总耗时秒 total_time_ms total_time * 1000 # 转换为毫秒更易读 return keypoints, total_time_ms # 返回关键点 耗时毫秒 # 主函数新增时间输出 if __name__ __main__: img_path /usr/test.jpg # 请确认图片存在 keypoints, infer_time_ms infer_keypoints(img_path) # 输出关键点 耗时 if keypoints: print(\n 关键点坐标x,y,置信度) joint_names [ 鼻子, 左眼, 右眼, 左耳, 右耳, 左肩, 右肩, 左肘, 右肘, 左手腕, 右手腕, 左髋, 右髋, 左膝, 右膝, 左踝, 右踝 ] for idx, (x, y, conf) in enumerate(keypoints): print(f {joint_names[idx]}: ({x}, {y}) 置信度{conf:.3f}) # 核心输出图片输入模型到2D坐标输出的时间 print(f\n⏱️ 图片输入模型 → 2D坐标输出 总耗时{infer_time_ms:.2f} 毫秒) else: print(❌ 关键点解析失败) print(f\n⏱️ 图片输入模型 → 2D坐标输出 总耗时{infer_time_ms:.2f} 毫秒) rknn_lite.release() print(\n✅ 资源释放完成)模型效果总结由于大家的在安装环境的过程会出现各种各样的问题所以大家在出现问题的时候多检查下自己的操作步骤多思考。最后祝愿大家都能够部署成功