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
运行与验证
- PC端执行
python convert_yolov5s.py,生成yolov5s.rknn,确认无报错且文件大小约15~20MB(i8量化后)。
- 将
rk3588_yolov5s.py、yolov5s.rknn和测试图片test.jpg(建议包含2~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偏移。