瑞芯微elf2开发板部署yolov5模型

发布于:2025-06-23 ⋅ 阅读:(15) ⋅ 点赞:(0)

随手记录一下,有时间了再完善

版本信息总结

组件 版本 说明
RKNN Lite 2.1.0 嵌入式端推理库
RKNN Runtime 2.1.0 运行时库 (967d001cc8)
RKNN Driver 0.9.8 NPU驱动程序
模型版本 6 RKNN模型格式版本
工具链版本 2.1.0+708089d1 模型转换工具链
Python 3.10 编程语言
OpenCV 4.x 图像处理库
目标平台 rk3588 Rockchip RK3588芯片

rklm-toolkit(瑞芯微模型工具 )、torch(深度学习框架 )、transformers(模型加载 ) 这些核心库,支撑模型转换、推理。

区分两个工具

  • rkllm-toolkit:更侧重 大语言模型(LLM) 相关的转换、部署(比如 DeepSeek 这类对话模型 )。
  • rknn-toolkit2:专门针对 计算机视觉模型(像 YOLOv5 ) 的转换、优化,适配 RK3588 的 NPU 加速。

运行脚本将onnx的模型转换一下

  • 关键参数(若需修改,可编辑 test.py):
    • INPUT_SIZE:模型输入尺寸(如 640,640)。
    • QUANTIZE_ON:是否开启量化(True 为 8 位量化,减小模型体积)。
    • model_path:ONNX 模型路径。
    • rknn_path:输出的 RKNN 模型路径(默认生成 yolov5s_relu.rknn)。

需要改一下转换脚本,因为默认是rk3566 ,与板子rk3588不适配。

修改好参数的脚本(适配RK3588)

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN

# Model from https://github.com/airockchip/rknn_model_zoo
ONNX_MODEL = 'yolov5s_relu.onnx'
RKNN_MODEL = 'yolov5s_relu_rk3588.rknn'  # 明确标识RK3588平台
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'

QUANTIZE_ON = True

OBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = 640

CLASSES = ("person", "bicycle", "car", "motorbike ", "aeroplane ", "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", "sofa",
           "pottedplant", "bed", "diningtable", "toilet ", "tvmonitor", "laptop	", "mouse	", "remote ", "keyboard ", "cell phone", "microwave ",
           "oven ", "toaster", "sink", "refrigerator ", "book", "clock", "vase", "scissors ", "teddy bear ", "hair drier", "toothbrush ")


def xywh2xyxy(x):
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2
    y[:, 1] = x[:, 1] - x[:, 3] / 2
    y[:, 2] = x[:, 0] + x[:, 2] / 2
    y[:, 3] = x[:, 1] + x[:, 3] / 2
    return y


def process(input, mask, anchors):
    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])

    box_confidence = input[..., 4]
    box_confidence = np.expand_dims(box_confidence, axis=-1)
    box_class_probs = input[..., 5:]
    box_xy = input[..., :2]*2 - 0.5

    col = np.tile(np.arange(0, grid_w), grid_w).reshape(-1, grid_w)
    row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_h)
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    grid = np.concatenate((col, row), axis=-1)
    box_xy += grid
    box_xy *= int(IMG_SIZE/grid_h)

    box_wh = pow(input[..., 2:4]*2, 2)
    box_wh = box_wh * anchors
    box = np.concatenate((box_xy, box_wh), axis=-1)
    return box, box_confidence, box_class_probs


def filter_boxes(boxes, box_confidences, box_class_probs):
    boxes = boxes.reshape(-1, 4)
    box_confidences = box_confidences.reshape(-1)
    box_class_probs = box_class_probs.reshape(-1, box_class_probs.shape[-1])

    _box_pos = np.where(box_confidences >= OBJ_THRESH)
    boxes = boxes[_box_pos]
    box_confidences = box_confidences[_box_pos]
    box_class_probs = box_class_probs[_box_pos]

    class_max_score = np.max(box_class_probs, axis=-1)
    classes = np.argmax(box_class_probs, axis=-1)
    _class_pos = np.where(class_max_score >= OBJ_THRESH)

    boxes = boxes[_class_pos]
    classes = classes[_class_pos]
    scores = (class_max_score * box_confidences)[_class_pos]
    return boxes, classes, scores


def nms_boxes(boxes, scores):
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]
    areas = w * h
    order = scores.argsort()[::-1]
    keep = []

    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep


def yolov5_post_process(input_data):
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
               [59, 119], [116, 90], [156, 198], [373, 326]]

    boxes, classes, scores = [], [], []
    for input, mask in zip(input_data, masks):
        b, c, s = process(input, mask, anchors)
        b, c, s = filter_boxes(b, c, s)
        boxes.append(b)
        classes.append(c)
        scores.append(s)

    boxes = np.concatenate(boxes)
    boxes = xywh2xyxy(boxes)
    classes = np.concatenate(classes)
    scores = np.concatenate(scores)

    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]
        keep = nms_boxes(b, s)
        nboxes.append(b[keep])
        nclasses.append(c[keep])
        nscores.append(s[keep])

    if not nclasses and not nscores:
        return None, None, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)
    return boxes, classes, scores


def draw(image, boxes, scores, classes):
    print("{:^12} {:^12}  {}".format('class', 'score', 'xmin, ymin, xmax, ymax'))
    print('-' * 50)
    for box, score, cl in zip(boxes, scores, classes):
        top, left, right, bottom = box
        top = int(top)
        left = int(left)
        right = int(right)
        bottom = int(bottom)

        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)

        print("{:^12} {:^12.3f} [{:>4}, {:>4}, {:>4}, {:>4}]".format(CLASSES[cl], score, top, left, right, bottom))

def letterbox(im, new_shape=(640, 640), color=(0, 0, 0)):
    shape = im.shape[:2]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    ratio = r, r
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
    dw /= 2
    dh /= 2

    if shape[::-1] != new_unpad:
        im = cv2.resize(im, 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))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
    return im, ratio, (dw, dh)


if __name__ == '__main__':
    rknn = RKNN(verbose=True)

    # 关键修改:config 阶段指定 target_platform 为 rk3588
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3588')
    print('done')

    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    print('--> Building model')
    # 移除 build 阶段的 target_platform 参数(避免版本不兼容)
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

    img = cv2.imread(IMG_PATH)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))

    print('--> Running model')
    img2 = np.expand_dims(img, 0)
    outputs = rknn.inference(inputs=[img2], data_format=['nhwc'])
    np.save('./onnx_yolov5_0.npy', outputs[0])
    np.save('./onnx_yolov5_1.npy', outputs[1])
    np.save('./onnx_yolov5_2.npy', outputs[2])
    print('done')

    input0_data = outputs[0]
    input1_data = outputs[1]
    input2_data = outputs[2]

    input0_data = input0_data.reshape([3, -1]+list(input0_data.shape[-2:]))
    input1_data = input1_data.reshape([3, -1]+list(input1_data.shape[-2:]))
    input2_data = input2_data.reshape([3, -1]+list(input2_data.shape[-2:]))

    input_data = list()
    input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
    input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))

    boxes, classes, scores = yolov5_post_process(input_data)

    img_1 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    if boxes is not None:
        draw(img_1, boxes, scores, classes)
        cv2.imwrite('result.jpg', img_1)
        print('Save results to result.jpg!')

    rknn.release()

python3 test.py

 保存好运行脚本即可,会在当前目录生成yolov5s_relu_rk3588.rknn

将转换好的yolov5s_relu_rk3588.rknn传到板端

(提前把需要用到的工具以及测试图片都传到板端目录里了)

检测脚本

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import cv2
import numpy as np
from rknnlite.api import RKNNLite
import os

# 全局参数
OBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = 640
CLASSES = ("person", "bicycle", "car", "motorbike ", "aeroplane ", "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", "sofa",
           "pottedplant", "bed", "diningtable", "toilet ", "tvmonitor", "laptop", "mouse", "remote ", "keyboard ", "cell phone", "microwave ",
           "oven ", "toaster", "sink", "refrigerator ", "book", "clock", "vase", "scissors ", "teddy bear ", "hair drier", "toothbrush ")

def load_rknn_model(model_path):
    """加载RKNN模型并初始化运行时"""
    rknn = RKNNLite()
    
    # 加载模型
    if rknn.load_rknn(model_path) != 0:
        print('模型加载失败')
        return None
    
    # 使用兼容性最好的初始化方法
    if rknn.init_runtime() != 0:
        print('运行环境初始化失败')
        return None
    
    print('RKNN Lite 初始化成功')
    return rknn

def preprocess(image):
    """图像预处理(与转换时完全一致)"""
    # BGR->RGB转换
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 直接resize到640x640
    image = cv2.resize(image, (IMG_SIZE, IMG_SIZE))
    # 添加batch维度 (NHWC格式)
    return np.expand_dims(image, axis=0)

def xywh2xyxy(x):
    """中心坐标转边界坐标"""
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # x1
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # y1
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # x2
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # y2
    return y

def process(input, mask, anchors):
    """处理单个输出层"""
    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])

    # 提取置信度和类别概率
    box_confidence = np.expand_dims(input[..., 4], axis=-1)
    box_class_probs = input[..., 5:]
    
    # 计算边界框中心
    box_xy = input[..., :2] * 2 - 0.5

    # 构建网格坐标
    col = np.tile(np.arange(grid_w), grid_h).reshape(grid_h, grid_w)
    row = np.tile(np.arange(grid_h).reshape(-1, 1), grid_w)
    grid = np.stack([col, row], axis=-1).reshape(grid_h, grid_w, 1, 2)
    
    # 调整边界框位置
    box_xy += grid
    box_xy *= int(IMG_SIZE / grid_h)
    
    # 计算边界框尺寸
    box_wh = pow(input[..., 2:4] * 2, 2) * anchors
    
    return np.concatenate((box_xy, box_wh), axis=-1), box_confidence, box_class_probs

def filter_boxes(boxes, confidences, class_probs):
    """过滤低置信度边界框"""
    boxes = boxes.reshape(-1, 4)
    confidences = confidences.reshape(-1)
    class_probs = class_probs.reshape(-1, class_probs.shape[-1])

    # 第一次过滤:目标置信度
    keep = confidences >= OBJ_THRESH
    boxes = boxes[keep]
    confidences = confidences[keep]
    class_probs = class_probs[keep]

    # 第二次过滤:类别置信度
    class_max = np.max(class_probs, axis=-1)
    classes = np.argmax(class_probs, axis=-1)
    keep = class_max >= OBJ_THRESH
    
    return boxes[keep], classes[keep], (class_max * confidences)[keep]

def nms_boxes(boxes, scores):
    """非极大值抑制"""
    x1, y1 = boxes[:, 0], boxes[:, 1]
    x2, y2 = boxes[:, 2], boxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)
    
    order = scores.argsort()[::-1]
    keep = []
    
    while order.size > 0:
        i = order[0]
        keep.append(i)
        
        # 计算IoU
        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:]])
        
        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)
        inter = w * h
        
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        
        # 保留低重叠框
        inds = np.where(iou <= NMS_THRESH)[0]
        order = order[inds + 1]
    
    return np.array(keep)

def yolov5_post_process(inputs):
    """YOLOv5后处理主函数"""
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
               [59, 119], [116, 90], [156, 198], [373, 326]]
    
    all_boxes, all_classes, all_scores = [], [], []
    
    # 处理三个输出层
    for input, mask in zip(inputs, masks):
        boxes, confs, probs = process(input, mask, anchors)
        boxes, classes, scores = filter_boxes(boxes, confs, probs)
        all_boxes.append(boxes)
        all_classes.append(classes)
        all_scores.append(scores)
    
    # 合并所有检测结果
    boxes = np.concatenate(all_boxes, axis=0)
    classes = np.concatenate(all_classes, axis=0)
    scores = np.concatenate(all_scores, axis=0)
    
    # 无检测结果处理
    if len(boxes) == 0:
        return None, None, None
    
    # 转换坐标格式
    boxes = xywh2xyxy(boxes)
    
    # 按类别进行NMS
    final_boxes, final_classes, final_scores = [], [], []
    for cls in set(classes):
        idx = classes == cls
        cls_boxes = boxes[idx]
        cls_scores = scores[idx]
        
        keep = nms_boxes(cls_boxes, cls_scores)
        final_boxes.append(cls_boxes[keep])
        final_classes.append(classes[idx][keep])
        final_scores.append(cls_scores[keep])
    
    return (np.concatenate(final_boxes),
            np.concatenate(final_classes),
            np.concatenate(final_scores))

def prepare_outputs(outputs):
    """对齐虚拟机输出处理逻辑"""
    out0 = outputs[0].reshape([3, -1] + list(outputs[0].shape[-2:]))
    out1 = outputs[1].reshape([3, -1] + list(outputs[1].shape[-2:]))
    out2 = outputs[2].reshape([3, -1] + list(outputs[2].shape[-2:]))
    
    return [
        np.transpose(out0, (2, 3, 0, 1)),
        np.transpose(out1, (2, 3, 0, 1)),
        np.transpose(out2, (2, 3, 0, 1))
    ]

def draw_results(image, boxes, classes, scores):
    """绘制检测结果"""
    for box, cls, score in zip(boxes, classes, scores):
        x1, y1, x2, y2 = map(int, box)
        label = f"{CLASSES[cls]} {score:.2f}"
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, label, (x1, y1-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    return image

if __name__ == "__main__":
    # 配置参数
    MODEL_PATH = "yolov5s_relu_rk3588.rknn"
    IMAGE_PATH = "person2.jpg"
    OUTPUT_PATH = "detection_result.jpg"
    
    # 加载模型
    print("="*50)
    print("RKNN YOLOv5 目标检测")
    print(f"模型: {MODEL_PATH}")
    print(f"图像: {IMAGE_PATH}")
    print("="*50)
    
    rknn = load_rknn_model(MODEL_PATH)
    if not rknn:
        exit(1)
    
    # 读取图像
    image = cv2.imread(IMAGE_PATH)
    if image is None:
        print(f"错误: 无法读取图像 {IMAGE_PATH}")
        rknn.release()
        exit(1)
    
    orig_h, orig_w = image.shape[:2]
    print(f"图像尺寸: {orig_w}x{orig_h}")
    
    # 预处理
    input_data = preprocess(image)
    
    # 推理
    outputs = rknn.inference(inputs=[input_data], data_format=["nhwc"])
    
    # 处理输出
    processed_outs = prepare_outputs(outputs)
    
    # 后处理
    boxes, classes, scores = yolov5_post_process(processed_outs)
    
    # 结果处理
    if boxes is not None:
        # 坐标映射回原图
        scale_x, scale_y = orig_w/IMG_SIZE, orig_h/IMG_SIZE
        boxes[:, [0, 2]] *= scale_x
        boxes[:, [1, 3]] *= scale_y
        
        # 绘制结果
        result_img = draw_results(image.copy(), boxes, classes, scores)
        cv2.imwrite(OUTPUT_PATH, result_img)
        
        # 打印结果
        print(f"检测到 {len(boxes)} 个目标:")
        for i, (box, cls, score) in enumerate(zip(boxes, classes, scores)):
            print(f"{i+1}. {CLASSES[cls]}: {score:.2f} @ [{int(box[0])},{int(box[1])} {int(box[2])},{int(box[3])}]")
        print(f"结果保存至: {OUTPUT_PATH}")
    else:
        print("未检测到目标")
    
    # 释放资源
    rknn.release()
    print("检测完成")

运行后会在当前目录生成detection_result.jpg

查看


网站公告

今日签到

点亮在社区的每一天
去签到