yolov5模型转换rknn

yolov5模型转换rknn

• • 安装virtualenv，用来来管理python环境
  • • 导出包到另一个系统
  • 图片标注
  • 安装yolo
  • onnx模型转换为rknn
  • RK3588安装RKNN-toolkit

安装virtualenv，用来来管理python环境

pip install virtualenvwrapper
mkdir ~/virtualenvs #创建存方环境的位置

查找virtualenvwrapper.sh脚本的位置

sudo find / -name virtualenvwrapper.sh
#/home/orangepi/.local/bin/virtualenvwrapper.sh

将virtualenvwrapper脚本的位置写入到.bashrc文件中，在.bashrc中添加

export WORKON_HOME=~/virtualenvs
source /home/orangepi/.local/bin/virtualenvwrapper.sh

运行： source ~/.bashrc

常用的一些命令

mkvirtualenv #创建环境
mkvirtualenv rknn --python=python3.9 #创建环境指定版本
workon  # 进入|切换 环境
lsvirtualenv # 展示环境列表
rmvirtualenv #删除环境
cpvirtualenv # 复制环境
deactivate # 退出当前环境

新建一个名为yolo的环境

mkvirtualenv yolo

导出包到另一个系统

输出虚拟环境中已安装包的名称及版本号到 requirements.txt 文件中：

pip freeze > requirements.txt

将文件发送到新设备

scp requirements.txt  orangepi@192.168.43.212:~/pi/rknn

新系统创建虚拟环境

mkvirtualenv rknn --python=python3.9 #创建虚拟环境

将 requirements.txt 包安装到虚拟环境里

pip install -r requirements.txt

图片标注

训练模型需要足够的目标照片，使用把视频提取一定数量的图片

import cv2cap = cv2.VideoCapture('/home/libai/orangePi/flightVideo/logs/a.avi')
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
imageNumber = 80 #截图张数
interval = round(total_frames / imageNumber) #间隔桢数
frameNum = 0 #当前桢数
cnt = 0 #图片数
path = '/home/libai/orangePi/rknn_yolo/myTestUav/images'
print(f'tatal frames:{total_frames},interval:{interval}')while True:success, image = cap.read()if frameNum % interval == 0 and success:cv2.imwrite(path + '%d.jpg' % cnt,image)cnt += 1print('>>output %d picture' % cnt)frameNum += 1if frameNum >= total_frames:print('ok')break
cap.release()

安装labelimg

pip install labelImg

新建文件夹目录，image/train存放图片，labels/train存放标注文件

打开labelimg
open dir 打开图片文件夹
change save dir 保存标签的路经
change save format 设置标签格式这里用yolo格式
A： 切换上一张照片
D：切换下一张照片
W：标注十字架

使用下面程序，将图片与标注文件按比例移动到val验证集文件夹
参考：https://blog.csdn.net/didiaopao/article/details/119927280?spm=1001.2014.3001.5501

import os, random, shutildef moveimg(fileDir, tarDir):pathDir = os.listdir(fileDir)  # 取图片的原始路径filenumber = len(pathDir)rate = 0.1  # 自定义抽取图片的比例，比方说100张抽10张，那就是0.1picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量图片sample = random.sample(pathDir, picknumber)  # 随机选取picknumber数量的样本图片print(sample)for name in sample:# print(fileDir + '/' + name, tarDir + "/" + name)shutil.move(fileDir + '/' + name, tarDir)returndef movelabel(file_list, file_label_train, file_label_val):for i in file_list:# filename = file_label_train + "\\" + i[:-4] + '.xml'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了filename = file_label_train + "/" + i[:-4] + '.txt'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了if os.path.exists(filename):shutil.move(filename, file_label_val)print(i + "处理成功！")if __name__ == '__main__':fileDir = r"/home/libai/orangePi/rknn_yolo/myTest/images/train"  # 源图片文件夹路径tarDir = r'/home/libai/orangePi/rknn_yolo/myTest/images/val'  # 图片移动到新的文件夹路径moveimg(fileDir, tarDir)file_list = os.listdir(tarDir)file_label_train = r"/home/libai/orangePi/rknn_yolo/myTest/labels/train"  # 源图片标签路径file_label_val = r"/home/libai/orangePi/rknn_yolo/myTest/labels/val"  # 标签# 移动到新的文件路径movelabel(file_list, file_label_train, file_label_val)

安装yolo

安装需求

git clone https://gitee.com/lmw0320/yolov5
cd yolov5
workon yolo #进入yolo虚拟环境
pip install -r requirements.txt

下载预训练权重文件，我用的是YOLOV5s.pt，放到yolov5/目录

在yolov5/data文件夹新建mytext.yaml 文件，写入下面内容，根据需要更改

train: /home/libai/orangePi/rknn_yolo/myTestPeople/images/train #训练图片位置
val: /home/libai/orangePi/rknn_yolo/myTestPeople/images/val#验证图片位置# Classes
nc: 1  # number of classes
names: ['object']  # class names

复制yolov5/models中yolov5s.yaml 重命名 mytextModel.yaml
将nc：80改为需要的数字，这里改为了1
修改train.py文件
–weights改为刚下载的yolov5s.pt位置
–cfg 改为刚新建的/models/mytestMode.yaml位置
–data 改为刚新建的/date/mytext.yaml位置

然后就可以训练自己的模型了

workon yolo
python train.py

使用tensorbord查看训练的参数

tensorboard --logdir=runs/train/xxx

本机打开http://localhost:6006/ 地址查看训练的参数

训练好后导出.rknn模型，best.pt 文件在./run/中

python export.py --include onnx --rknpu RK3588S --weights ./best.pt

onnx模型转换为rknn

下载RKNN-Toolkit2 工具 https://eyun.baidu.com/s/3eTDMk6Y密码rknn
新建rknn环境，安装依赖

mkvirtualenv rknn
cd rknn-toolkit2-1.4.0
pip3 install -r doc/requirements*.txt

安装RKNN-Toolkit2

cd package/
sudo pip3 install rknn_toolkit2-1.4.0_22dcfef4-cp38-cp38-linux_x86_64.whl

查看是否安装好了

from rknn.api import RKNN

没报错
修改rknn-toolkit2-1.4.0/examples/onnx/yolov5/text.py 文件，ONNX_MODEL为刚生成的.onnx文件，RKNN_MODEL为将要转换成.rknn的文件

ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'best.rknn'
IMG_PATH = './images4.jpg'
DATASET = './dataset.txt'

需要修改一个位置

rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]],target_platform='rk3588') 

然后运行程序就可以看到结果了

RK3588安装RKNN-toolkit

RKNN Toolkit Lite2 v1.4支持运行于 Debian 10 / 11 (aarch64) 操作系统支持python3.7 / 3.9

安装python3.9

sudo apt install python3.9
mkvirtualenv rknn --python=python3.9 #创建虚拟环境

安装 RKNN Toolkit Lite2

pip3 install rknn_toolkit_lite2-1.x.0-cp39-cp39m-linux_aarch64.whl

报错了

尝试安装python3.9-dev

 sudo  apt-get install python3.9-dev

import RKNNLite 如果没有报错就安装好了

import RKNNLite

用下面板载的程序，这个程序修改于rknn-toolkit2-1.4.0中的examples/onnx/yolov5/text.，测试一下,使用刚转换的best.rknn模型，推理名为images4.jpg的照片

import numpy as np
import cv2
from rknnlite.api import RKNNLite
import timeclass YOLOV5():def __init__(self):self.RKNN_MODEL = './best.rknn'self.IMG_PATH = './images4.jpg'self.OBJ_THRESH = 0.25self.NMS_THRESH = 0.45self.IMG_SIZE = 640self.CLASSES = ("person")self.yoloInit()def yoloInit(self):# Create RKNN objectself.rknn = RKNNLite()# load RKNN modelprint('--> Load RKNN model')ret = self.rknn.load_rknn(self.RKNN_MODEL)# Init runtime environmentprint('--> Init runtime environment')ret = self.rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)  #使用0 1 2三个NPU核心if ret != 0:print('Init runtime environment failed!')exit(ret)print('done')def inference(self):# Set inputsimg = cv2.imread(self.IMG_PATH)img = cv2.copyMakeBorder(img, 0, 160, 0, 0, cv2.BORDER_CONSTANT, value=(0,0,0))  # add borderimg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)# Inferencea = time.time()outputs = self.rknn.inference(inputs=[img])b = time.time()print('time:',b-a)# post processinput0_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 = self.yolov5_post_process(input_data)print(boxes,classes,scores)self.rknn.release()def sigmoid(self,x):return 1 / (1 + np.exp(-x))def xywh2xyxy(self,x):# Convert [x, y, w, h] to [x1, y1, x2, y2]y = np.copy(x)y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left xy[:, 1] = x[:, 1] - x[:, 3] / 2  # top left yy[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right xy[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right yreturn ydef process(self,input, mask, anchors):anchors = [anchors[i] for i in mask]grid_h, grid_w = map(int, input.shape[0:2])box_confidence = self.sigmoid(input[..., 4])box_confidence = np.expand_dims(box_confidence, axis=-1)box_class_probs = self.sigmoid(input[..., 5:])box_xy = self.sigmoid(input[..., :2])*2 - 0.5col = 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 += gridbox_xy *= int(self.IMG_SIZE/grid_h)box_wh = pow(self.sigmoid(input[..., 2:4])*2, 2)box_wh = box_wh * anchorsbox = np.concatenate((box_xy, box_wh), axis=-1)return box, box_confidence, box_class_probsdef filter_boxes(self,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 >= self.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 >= self.OBJ_THRESH)boxes = boxes[_class_pos]classes = classes[_class_pos]scores = (class_max_score* box_confidences)[_class_pos]return boxes, classes, scoresdef nms_boxes(self,boxes, scores):x = boxes[:, 0]y = boxes[:, 1]w = boxes[:, 2] - boxes[:, 0]h = boxes[:, 3] - boxes[:, 1]areas = w * horder = 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 * h1ovr = inter / (areas[i] + areas[order[1:]] - inter)inds = np.where(ovr <= self.NMS_THRESH)[0]order = order[inds + 1]keep = np.array(keep)return keepdef yolov5_post_process(self,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 = self.process(input, mask, anchors)b, c, s = self.filter_boxes(b, c, s)boxes.append(b)classes.append(c)scores.append(s)boxes = np.concatenate(boxes)boxes = self.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 = self.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, Noneboxes = np.concatenate(nboxes)classes = np.concatenate(nclasses)scores = np.concatenate(nscores)return boxes, classes, scoresyolo = YOLOV5()
yolo.inference()

推理时间大概0.018秒，此前在电脑上推理时间为1.2秒，说明板载的npu加速还是很厉害的
![在这里插入图片描述](https://img-blog.csdnimg.cn/5ecfc8a0fd56469da7855a3f2378ea08.png
现在就可以使用这块芯片干更有意思的事情了