代码收藏家 基于RK3588 板子实现yolov5 模型转化推理—-实测python 版
1. yolov5 模型训练
模型训练可以基于ubuntu or windowns
1.1 下载yolov5模型
YOLOV5https://github.com/ultralytics/yolov5
1.2 模型训练
可以参考官方文档，这里就不详细介绍了。将训练好的权重文件备用
2. Pt —>ONNX 模型转化
2.1在训练好的pt 模型后，找到yolov5 文件下export.py 文件
2.2 修改export 下模型文件，配置文件等数据
修改 “opset” 修改onnx 的版本，要改为11or 12
如果没有安装onnx 则执行
pip install onnx==1.12最后修改转化文件的类型，onnx
在找到 models/yolo.py 文件
将原本的 forward(self,x ) 函数用一下代码替代：注意，仅在模型转化时才替换一下该函数，若执行模型训练和推理时，记得要修改会原来代码。
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
z.append(torch.sigmoid(self.m[i](x[i]))) # conv output
return z
最后执行export.py 文件，生成onnx 模型结果如下：
3.ONNX—>RKNN 模型转化
onnx 模型转为rknn 需要在ubuntu 环境下
3.1 rknn 官方代码下载
RKNN_YOLOV5_DEMOhttps://github.com/rockchip-linux/rknn-toolkit2/tree/master/rknpu2/examples/rknn_yolov5_demo
1.1安装环境
创建虚拟环境，创建一个新的 conda 环境，python版本设置为3.6、3.8或者3.10
conda create -n rknn python=3.8激活环境
conda activate rknn安装rknn 工具包：下载rknn 工具包
$ pip install rknn_toolkit2-1.5.0+1fa95b5c-cp38-cp38-linux_x86_64.whl确认安装成功 检查 RKNN-Toolkit2 是否安装成功
(rknn)$ python
(rknn)$ from rknn.api import RKNN如未报错说明安装成功
3.2 修改onnx2rknn.py
platform = 'rk3588'
exp = 'face_best'
Width = 640
Height = 640
# Model from https://github.com/airockchip/rknn_model_zoo
MODEL_PATH = '/rk3588/rknpu2/examples/rknn_yolov5_demo/convert_rknn_demo/yolov5/onnx_models/face-1best.onnx'
NEED_BUILD_MODEL = True
# NEED_BUILD_MODEL = False
im_file = '/ydemo_ai/rknpu2/examples/rknn_yolov5_demo/convert_rknn_demo/yolov5/frame_1.jpg'3.3 执行onnx2rknn.py
4.RKNN 模型推理
4.1将生成的rknn 模型移到板子中，板子中的同样安装rknn 工具包
rknn_toolkit_lite2https://github.com/rockchip-linux/rknn-toolkit2/tree/master/rknn_toolkit_lite2找到 rknn_toollkit_lite2/packages ，里面有所需环境安装包， 上文在ubuntu下安装的是 rknn-toolkit 注意区分，一个是ubuntu 工具环境, lite2 为 rknn 板端环境。
创建虚拟环境，记得使用miconda 安装教程可自行搜索
在虚拟环境下安装;
pip install rknn_toolkit_lite2-2.3.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl4.2 推理代码 demo
修改模型地址和测试文件即可，支持图像，视频.
from copy import copy
import time
import numpy as np
import cv2
from rknnlite.api import RKNNLite
RKNN_MODEL = '/home/linaro/rknn-toolkit-lite2/onnx2rknn/face_best-640-640_rk3588.rknn'
# IMG_PATH = './input/240513_00000741.jpg'
OBJ_THRESH = 0.25
NMS_THRESH = 0.45
IMG_SIZE = (640, 640)
OUTPUT_VIDEO_PATH = 'output_1.mp4'
BOX = (0, 0, 640, 640)
# CLASSES = ['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']
# CLASSES = ['tv', 'phone', 'battery', 'clothes', 'seat', 'people', 'Keyboard']
CLASSES = ['PERSON']
anchors = [[[10, 13], [16, 30], [33, 23]],
[[30, 61], [62, 45], [59, 119]],
[[116, 90], [156, 198], [373, 326]]]
class Letter_Box_Info():
def __init__(self, shape, new_shape, w_ratio, h_ratio, dw, dh, pad_color) -> None:
self.origin_shape = shape
self.new_shape = new_shape
self.w_ratio = w_ratio
self.h_ratio = h_ratio
self.dw = dw
self.dh = dh
self.pad_color = pad_color
def box_process(position, anchors):
grid_h, grid_w = position.shape[2:4]
col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h)) # (80, 80) (80, 80)
col = col.reshape(1, 1, grid_h, grid_w) # (1, 1, 80, 80)
row = row.reshape(1, 1, grid_h, grid_w)
grid = np.concatenate((col, row), axis=1) # (1, 2, 80, 80)
stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1) # 8 8
col = col.repeat(len(anchors), axis=0)
row = row.repeat(len(anchors), axis=0)
anchors = np.array(anchors)
anchors = anchors.reshape(*anchors.shape, 1, 1) # (3, 2, 1, 1)
box_xy = position[:,:2,:,:]*2 - 0.5
box_wh = pow(position[:,2:4,:,:]*2, 2) * anchors
box_xy += grid
box_xy *= stride
box = np.concatenate((box_xy, box_wh), axis=1) # (3, 4, 80, 80)
# Convert [c_x, c_y, w, h] to [x1, y1, x2, y2]
xyxy = np.copy(box)
xyxy[:, 0, :, :] = box[:, 0, :, :] - box[:, 2, :, :]/ 2 # top left x
xyxy[:, 1, :, :] = box[:, 1, :, :] - box[:, 3, :, :]/ 2 # top left y
xyxy[:, 2, :, :] = box[:, 0, :, :] + box[:, 2, :, :]/ 2 # bottom right x
xyxy[:, 3, :, :] = box[:, 1, :, :] + box[:, 3, :, :]/ 2 # bottom right y
return xyxy
#
def filter_boxes(boxes, box_confidences, box_class_probs):
"""Filter boxes with object threshold.
"""
box_confidences = box_confidences.reshape(-1)
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* box_confidences >= OBJ_THRESH)
scores = (class_max_score* box_confidences)[_class_pos]
boxes = boxes[_class_pos]
classes = classes[_class_pos]
return boxes, classes, scores
# def filter_boxes(boxes, box_confidences, box_class_probs):
# """Filter boxes with object threshold.
# """
# 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):
"""Suppress non-maximal boxes.
# Returns
keep: ndarray, index of effective boxes.
"""
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 post_process(input_data, anchors):
boxes, scores, classes_conf = [], [], []
# 1*255*h*w -> 3*85*h*w
input_data = [_in.reshape([len(anchors[0]),-1]+list(_in.shape[-2:])) for _in in input_data]
for i in range(len(input_data)): # (3, 85, 80, 80)
boxes.append(box_process(input_data[i][:,:4,:,:], anchors[i])) # (3, 4, 80, 80)
scores.append(input_data[i][:,4:5,:,:]) # (3, 1, 80, 80)
classes_conf.append(input_data[i][:,5:,:,:]) # (3, 80, 80, 80)
def sp_flatten(_in):
ch = _in.shape[1]
_in = _in.transpose(0,2,3,1)
return _in.reshape(-1, ch)
boxes = [sp_flatten(_v) for _v in boxes] # (3, 19200, 4)
classes_conf = [sp_flatten(_v) for _v in classes_conf] # (3, 19200, 80)
scores = [sp_flatten(_v) for _v in scores] # (3, 19200, 1)
boxes = np.concatenate(boxes) # (25200, 4)
classes_conf = np.concatenate(classes_conf) # (25200, 80)
scores = np.concatenate(scores) # (25200, 1)
# filter according to threshold
boxes, classes, scores = filter_boxes(boxes, scores, classes_conf)
# (12, 4) 12 12
# nms
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)
if len(keep) != 0:
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):
for box, score, cl in zip(boxes, scores, classes):
print("cl:",cl)
top, left, right, bottom = [int(_b) for _b in box]
print("%s @ (%d %d %d %d) %.3f" % (CLASSES[cl], top, left, right, bottom, score))
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)
def letterbox(im, new_shape=(640, 640), color=(0, 0, 0), letter_box_info_list=[]):
shape = im.shape[:2] # current shape [height, width]
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 # width, height ratios
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] # wh padding
# dw, dh = np.mod(dw, 32), np.mod(dh, 32)
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
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) # add border
letter_box_info_list.append(Letter_Box_Info(shape, new_shape, ratio, ratio, dw, dh, color))
return im, letter_box_info_list
def get_real_box(box, in_format='xyxy', letter_box_info_list=[]):
bbox = copy(box)
# unletter_box result
if in_format=='xyxy':
bbox[:,0] -= letter_box_info_list[-1].dw
bbox[:,0] /= letter_box_info_list[-1].w_ratio
bbox[:,0] = np.clip(bbox[:,0], 0, letter_box_info_list[-1].origin_shape[1])
bbox[:,1] -= letter_box_info_list[-1].dh
bbox[:,1] /= letter_box_info_list[-1].h_ratio
bbox[:,1] = np.clip(bbox[:,1], 0, letter_box_info_list[-1].origin_shape[0])
bbox[:,2] -= letter_box_info_list[-1].dw
bbox[:,2] /= letter_box_info_list[-1].w_ratio
bbox[:,2] = np.clip(bbox[:,2], 0, letter_box_info_list[-1].origin_shape[1])
bbox[:,3] -= letter_box_info_list[-1].dh
bbox[:,3] /= letter_box_info_list[-1].h_ratio
bbox[:,3] = np.clip(bbox[:,3], 0, letter_box_info_list[-1].origin_shape[0])
return bbox
if __name__ == '__main__':
def sigmoid(x):
return 1 / (1 + np.exp(-x))
rknn = RKNNLite()
print('--> Load RKNN model')
ret = rknn.load_rknn(RKNN_MODEL)
if ret != 0:
print('Load RKNN model failed')
exit(ret)
print('done')
ret = rknn.init_runtime()
if ret != 0:
print('Init runtime environment failed!')
exit(ret)
print('done')
cap = cv2.VideoCapture(0)
# # 获取视频的一些属性
# fps = cap.get(cv2.CAP_PROP_FPS)
# width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
# height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# 创建 VideoWriter 对象
# x, y, w, h = BOX
# fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 或者使用 'XVID'
# out = cv2.VideoWriter(OUTPUT_VIDEO_PATH, fourcc, fps, (w, h))
# fps = 0.0
# while True:
t1 = time.time()
# 读取一帧
# ret, frame1 = cap.read()
# if not ret:
# break
# img0 = frame1
img0 = cv2.imread('/home/linaro/rknn-toolkit-lite2/predict_img/office_data.jpg')
# img0 = img0 / 255.0
# img0 = img0.astype(np.int8)
# 加载帧
# img0 = frame
img, letter_box_info_list = letterbox(im= img0.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0])) # padded resize
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # HWC to CHW, BGR to RGB
# img = np.expand_dims(img, axis=0)
# img = img.astype(np.int8)
if len(img.shape) == 3:
img = img[None] # expand for batch dim
outputs = rknn.inference(inputs=[img]) # Inference
boxes, classes, scores = post_process(outputs, anchors)
print("classes",classes)
print("scores",scores)
boxes_filter, scores_filter, classes_filter = [0, 0, 0, 0], [], []
# scores_filter = np.array(scores)
# classes_filter = np.array(classes)
# processed_scores = sigmoid(scores_filter)
# processed_classes = sigmoid(classes_filter)
# print("scores_filter",processed_scores)
# print("classes_filter",processed_classes)
mask = scores > 0.4
scores_filter = scores[mask]
boxes_filter = boxes[mask]
classes_filter = classes[mask]
# 筛选 scores 和对应的 boxes
# scores_filter = processed_scores[mask]
# boxes_filter = boxes[mask]
# classes_filter = processed_classes[mask]
# max_box = [0, 0, 0, 0]
# for box, score, cl in zip(boxes, scores, classes):
# print("cl",cl)
# # if cl == 0:
# # if (box[2]-box[0])*(box[3]-box[1]) > (max_box[2]-max_box[0])*(max_box[3]-max_box[1]):
# max_box = box
# boxes_filter = np.expand_dims(max_box, axis=0)
# scores_filter = np.expand_dims(score, axis=0)
# classes_filter = np.expand_dims(cl, axis=0)
# img_p = img0.copy()
# print("boxes_filter",boxes_filter)
draw(img0, get_real_box(boxes_filter, 'xyxy', letter_box_info_list), scores_filter, classes_filter)
print("run time:",time.time()-t1)
# cv2.namedWindow('Window', cv2.WINDOW_NORMAL)
# cv2.imshow("Window",img0)
# cv2.waitKey(1)
cv2.imwrite(f"./res/{time.time()}.jpg", img0)
print("run time:",time.time()-t1)
# out.write(img_p)
rknn.release()
推理结果：
作者：码农的日常搅屎棍
