4. SSD-MobileNetV2目标检测部署

任务目标

将SSD-MobileNetV2（COCO预训练，80类）转换为RKNN格式，在RK3588上完成目标检测推理，输出带标注的检测结果图，并统计单帧推理耗时。

环境与依赖

  • PC端：Ubuntu 20.04+，rknn-toolkit2（≥1.6.0），onnx==1.12.0
  • 板端：RK3588，RKNN-Toolkit-Lite2，Python 3.8+，OpenCV
  • 模型来源：ssd_mobilenet_v2_coco.onnx（可从TensorFlow Model Zoo转换或公开ONNX仓库获取，输入尺寸300×300）

关键实现

以下代码在PC端完成ONNX→RKNN转换与量化，并生成可在板端直接运行的推理脚本。核心处理链为：图像预处理（letterbox缩放至300×300）→ RKNN推理 → 后处理（解码SSD输出锚框、NMS过滤）→ 绘制结果。

# convert_ssd_mobilenet.py（PC端执行）
from rknn.api import RKNN

ONNX_MODEL = 'ssd_mobilenet_v2_coco.onnx'
RKNN_MODEL = 'ssd_mobilenet_v2_coco.rknn'
DATASET = 'dataset.txt'  # 每行一张量化校准图片路径，建议20~50张

rknn = RKNN()
# 配置：SSD需指定reorder_channel为'0 1 2'（RGB顺序）
rknn.config(mean_values=[[127.5, 127.5, 127.5]], std_values=[[127.5, 127.5, 127.5]], reorder_channel='0 1 2', target_platform='rk3588')
# 加载ONNX模型
ret = rknn.load_onnx(model=ONNX_MODEL)
assert ret == 0, '模型加载失败'
# 构建RKNN（若需量化，传入dataset.txt；否则do_quantization=False）
ret = rknn.build(do_quantization=True, dataset=DATASET)
assert ret == 0, '构建失败'
# 导出
ret = rknn.export_rknn(RKNN_MODEL)
assert ret == 0, '导出失败'
print('转换完成:', RKNN_MODEL)

# rknn_ssd_detect.py（板端执行）
import cv2
import numpy as np
from rknnlite.api import RKNNLite

RKNN_MODEL = 'ssd_mobilenet_v2_coco.rknn'
IMG_PATH = 'test.jpg'
OUT_PATH = 'result.jpg'
CONF_THRESH = 0.5
NMS_THRESH = 0.45
NUM_CLASSES = 90  # COCO类别数（含背景）
NUM_ANCHORS = 1917  # SSD-MobileNetV2默认锚框总数

# COCO 80类名称（索引1~80对应实际类别，0为背景）
COCO_LABELS = ['background','person','bicycle','car','motorcycle','airplane','bus','train','truck','boat',
               'traffic light','fire hydrant','stop sign','parking meter','bench','bird','cat','dog','horse','sheep',
               'cow','elephant','bear','zebra','giraffe','backpack','umbrella','handbag','tie','suitcase',
               'frisbee','skis','snowboard','sports ball','kite','baseball bat','baseball glove','skateboard','surfboard','tennis racket',
               'bottle','wine glass','cup','fork','knife','spoon','bowl','banana','apple','sandwich',
               'orange','broccoli','carrot','hot dog','pizza','donut','cake','chair','couch','potted plant',
               'bed','dining table','toilet','tv','laptop','mouse','remote','keyboard','cell phone','microwave',
               'oven','toaster','sink','refrigerator','book','clock','vase','scissors','teddy bear','hair drier',
               'toothbrush']

def letterbox(img, size=300):
    h, w = img.shape[:2]
    scale = min(size / w, size / h)
    nw, nh = int(w * scale), int(h * scale)
    resized = cv2.resize(img, (nw, nh))
    canvas = np.full((size, size, 3), 114, dtype=np.uint8)
    x, y = (size - nw) // 2, (size - nh) // 2
    canvas[y:y+nh, x:x+nw] = resized
    return canvas, scale, x, y

def decode_outputs(outputs, scale, pad_x, pad_y):
    # outputs: [1, 1917, 4+1+90] 即 [boxes, conf, class_scores]
    preds = outputs[0][0]  # (1917, 95)
    boxes = preds[:, :4]   # (cx, cy, w, h) 归一化坐标
    scores = preds[:, 4:5] * preds[:, 5:]  # conf * class_prob
    class_ids = np.argmax(scores, axis=1)
    confs = np.max(scores, axis=1)
    # 还原到原图坐标
    boxes[:, 0] = (boxes[:, 0] * 300 - pad_x) / scale  # x_min
    boxes[:, 1] = (boxes[:, 1] * 300 - pad_y) / scale  # y_min
    boxes[:, 2] = (boxes[:, 2] * 300) / scale          # w
    boxes[:, 3] = (boxes[:, 3] * 300) / scale          # h
    boxes[:, 2] += boxes[:, 0]  # x_max
    boxes[:, 3] += boxes[:, 1]  # y_max
    return boxes, class_ids, confs

def nms(boxes, scores, thresh):
    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        order = order[np.where(iou <= thresh)[0] + 1]
    return keep

# 初始化NPU
rknn = RKNNLite()
ret = rknn.load_rknn(RKNN_MODEL)
assert ret == 0, 'RKNN加载失败'
ret = rknn.init_runtime()
assert ret == 0, 'NPU初始化失败'

# 推理
img = cv2.imread(IMG_PATH)
canvas, scale, pad_x, pad_y = letterbox(img)
input_tensor = canvas[:, :, ::-1]  # BGR→RGB
outputs = rknn.inference(inputs=[input_tensor])
boxes, class_ids, confs = decode_outputs(outputs, scale, pad_x, pad_y)

# 过滤低置信度
mask = confs > CONF_THRESH
boxes, class_ids, confs = boxes[mask], class_ids[mask], confs[mask]
if len(boxes) > 0:
    keep = nms(boxes, confs, NMS_THRESH)
    boxes, class_ids, confs = boxes[keep], class_ids[keep], confs[keep]

# 绘制结果
for box, cls, conf in zip(boxes, class_ids, confs):
    if cls < 1 or cls > 80:
        continue
    x1, y1, x2, y2 = box.astype(int)
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    label = f'{COCO_LABELS[cls]} {conf:.2f}'
    cv2.putText(img, label, (x1, max(0, y1-5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
cv2.imwrite(OUT_PATH, img)
print(f'检测完成，共{len(boxes)}个目标，结果保存至{OUT_PATH}')
rknn.release()

运行与验证

  1. PC端转换：执行python convert_ssd_mobilenet.py，生成ssd_mobilenet_v2_coco.rknn。若需量化，确保dataset.txt中的图片路径有效。
  2. 部署到板端：将.rknn文件、推理脚本和测试图片拷贝至RK3588。
  3. 板端推理：执行python rknn_ssd_detect.py。
  4. 通过标准：程序正常结束并输出检测目标数量；打开result.jpg，图中目标被绿色框正确标注且类别置信度合理（如人、车、狗等常见物体）。

性能评估方法

在推理脚本中，对同一张图连续执行50次rknn.inference()，取平均耗时作为单帧推理时间。注意首次推理包含NPU初始化，不计入统计。典型性能范围：RK3588 NPU上SSD-MobileNetV2 INT8量化模型单帧推理约10~20ms（实际值取决于NPU频率与内存带宽）。

常见问题

  • 输出维度不匹配：确认ONNX模型输出是否为[1, 1917, 95]格式。若输出为多个分支（如SSD原始输出），需在转换前用ONNX简化脚本合并为单输出。
  • 检测框偏移：检查letterbox的填充坐标还原是否正确，以及是否使用了RGB输入（RKNN默认RGB顺序）。
  • 量化精度下降：增加dataset.txt中校准图片数量至100张以上，或改用do_quantization=False进行FP16推理对比。
← 上一章 📖 返回目录 下一章 →