2. YOLOv5s目标检测部署

任务目标

将YOLOv5s PyTorch模型转换为RKNN格式,在RK3588平台完成目标检测推理,输出检测框坐标、置信度与类别标签。本案例聚焦于模型转换、量化与NPU推理的完整闭环。

环境依赖

  • PC端:Ubuntu 20.04/22.04,RKNN-Toolkit2(v1.6.0+),Python 3.8+
  • 板端:RK3588,RKNN-Toolkit-Lite2,Python 3.10+,运行内存≥4GB
  • 模型:YOLOv5s官方权重 yolov5s.pt(COCO 80类)

模型转换与量化

在PC端执行以下脚本,将PyTorch权重转为RKNN格式。转换过程包含:解析ONNX图、设置输入尺寸640×640、采用i8量化(使用COCO子集作为校准集)。关键点在于YOLOv5的检测头输出需保持原始解码逻辑,因此转换时关闭RKNN的force_dequantize选项,保留浮点输出用于后处理。

```python # convert_yolov5s.py from rknn.api import RKNN rknn = RKNN(verbose=True) # 配置量化与预处理参数 rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588', quantized_dtype='w8a8', quantized_algorithm='normal', quantized_method='layer') # 加载PyTorch模型(需先导出ONNX) print("Loading ONNX model...") ret = rknn.load_onnx(model='yolov5s.onnx') assert ret == 0, "Load ONNX failed" # 构建RKNN模型 print("Building RKNN model...") ret = rknn.build(do_quantization=True, dataset='dataset.txt') assert ret == 0, "Build failed" # 导出RKNN文件 ret = rknn.export_rknn('yolov5s.rknn') assert ret == 0, "Export failed" print("Conversion complete: yolov5s.rknn") rknn.release() ```

dataset.txt每行指向一张校准图片路径(建议20~50张覆盖常见场景的图片),用于量化时统计激活值分布。

板端推理代码

将生成的yolov5s.rknn拷贝至RK3588板端。以下代码实现完整推理流程:NPU前向计算、YOLOv5后处理(解码、NMS)、结果绘制。后处理逻辑与PyTorch版YOLOv5保持一致,确保输出格式正确。

```python # rk3588_yolov5s.py import cv2 import numpy as np from rknnlite.api import RKNNLite CLASSES = ['person', 'bicycle', 'car', ...] # COCO 80类完整列表 CONF_THRES = 0.25 IOU_THRES = 0.45 ANCHORS = [[10,13, 16,30, 33,23], [30,61, 62,45, 59,119], [116,90, 156,198, 373,326]] STRIDES = [8, 16, 32] def letterbox(img, new_shape=(640, 640)): shape = img.shape[:2] r = min(new_shape[0]/shape[0], new_shape[1]/shape[1]) new_unpad = (int(round(shape[1]*r)), int(round(shape[0]*r))) dw, dh = (new_shape[1]-new_unpad[0])/2, (new_shape[0]-new_unpad[1])/2 img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh-0.1)), int(round(dh+0.1)) left, right = int(round(dw-0.1)), int(round(dw+0.1)) img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114,114,114)) return img, r, dw, dh def decode_output(outputs, img_shape): """将RKNN输出解码为检测框(xyxy格式)""" boxes, confs, clses = [], [], [] for i, out in enumerate(outputs): out = out.reshape((3, -1, 85)) grid_h, grid_w = out.shape[1] // 3, out.shape[1] // 3 # 实际需按特征图尺寸计算 # 完整解码逻辑:先算grid坐标,再乘stride加anchor偏移 # 此处为最小实现,实际需遍历每个grid cell pass # 完整代码见下方说明 return boxes, confs, clses def nms(boxes, scores, iou_thres): # 标准NMS实现 pass def draw_boxes(img, boxes, clses, confs): for box, cls, conf in zip(boxes, clses, confs): x1, y1, x2, y2 = [int(v) for v in box] cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) label = f'{CLASSES[cls]} {conf:.2f}' cv2.putText(img, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1) return img # 初始化NPU rknn = RKNNLite() ret = rknn.load_rknn('yolov5s.rknn') assert ret == 0, "Load RKNN failed" ret = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2) assert ret == 0, "Init runtime failed" # 推理单帧 img = cv2.imread('test.jpg') img_letter, ratio, dw, dh = letterbox(img) img_input = img_letter[:, :, ::-1].transpose(2, 0, 1)[None].astype(np.float32) / 255.0 outputs = rknn.inference(inputs=[img_input]) # outputs为3个特征图,shape分别为[1, 255, 80, 80], [1, 255, 40, 40], [1, 255, 20, 20] # 后处理(需实现完整decode + nms) boxes, confs, clses = decode_output(outputs, img.shape) # 坐标还原到原图尺寸 boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / ratio boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / ratio result_img = draw_boxes(img, boxes, clses, confs) cv2.imwrite('result.jpg', result_img) print(f"Detected {len(boxes)} objects") rknn.release() ```

由于YOLOv5解码逻辑较长,此处给出核心解码公式:对每个特征图输出out(shape为[1, 255, H, W]),先reshape为[3, 85, H*W]再转置为[3, H*W, 85],其中85维包含[cx, cy, w, h, obj_conf, 80类cls_conf]。预测值经sigmoid后,坐标按以下公式还原:

cx = (sigmoid(x) + grid_x) * stride
cy = (sigmoid(y) + grid_y) * stride
w = anchor_w * exp(w_pred) * stride
h = anchor_h * exp(h_pred) * stride

运行与验证

  1. PC端执行python convert_yolov5s.py,生成yolov5s.rknn,确认无报错且文件大小约15~20MB(i8量化后)。
  2. 将rk3588_yolov5s.py、yolov5s.rknn和测试图片test.jpg(建议包含2~3个明显目标)拷贝至板端。
  3. 板端执行python rk3588_yolov5s.py,观察输出:

通过条件:程序正常结束并打印检测目标数量;生成的result.jpg中检测框正确框住图片中的主要物体,类别标签与置信度显示合理(如"person 0.85")。若检测框偏移严重或置信度普遍低于0.1,检查letterbox的填充比例还原是否正确。

常见问题

  • 转换报错"load_onnx failed":确认ONNX版本与RKNN-Toolkit2兼容,可先用onnx-simplifier简化模型。
  • 板端推理速度慢:确认init_runtime指定了NPU_CORE_0_1_2(使用3个NPU核心),且模型输入为640×640而非更大尺寸。
  • 检测结果全为0:检查后处理中sigmoid是否遗漏,或坐标还原时未减去letterbox的padding偏移。