目录
这里直接参考我的另一篇文章:yolov8训练pt模型转换为rknn模型_部署在RK3588上--整个流程-CSDN博客
ls /dev/video*
拍摄照片收集数据集
python capture_image5.py
import cv2 as cv
import os
import datetime
import numpy as np # 确保这一行在文件顶部
def create_directory_if_not_exists(directory):
if not os.path.exists(directory):
os.makedirs(directory)
# 鼠标点击事件的回调函数
def mouse_click(event, x, y, flags, param):
global frame, image_dir, save_count
if event == cv.EVENT_LBUTTONDOWN:
# 当鼠标左键点击时保存图片,使用PNG格式保存以确保无损
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
image_filename = os.path.join(image_dir, f"image_{timestamp}.png")
cv.imwrite(image_filename, frame) # 使用默认参数保存PNG,确保无损
print(f"图片已保存为 {image_filename}")
save_count += 1
print(f"已保存图片总数: {save_count}")
def capture_image():
global frame, image_dir, save_count
# 初始化保存计数器
save_count = 0
# 创建存放图片的目录
image_dir = 'images_1'
create_directory_if_not_exists(image_dir)
# 尝试打开默认摄像头
cap = cv.VideoCapture(0) # 根据实际情况选择摄像头编号
if not cap.isOpened():
print("无法打开摄像头")
return
# 设置摄像头属性
cap.set(cv.CAP_PROP_FRAME_WIDTH, 1920)
cap.set(cv.CAP_PROP_FRAME_HEIGHT, 1080)
cap.set(cv.CAP_PROP_FPS, 30)
cap.set(cv.CAP_PROP_FOURCC, cv.VideoWriter_fourcc('M', 'J', 'P', 'G'))
# 尝试调整图像质量相关的属性
cap.set(cv.CAP_PROP_EXPOSURE, -6) # 调整曝光时间
cap.set(cv.CAP_PROP_GAIN, 0) # 固定增益
cap.set(cv.CAP_PROP_WHITE_BALANCE_BLUE_U, 5000) # 白平衡
cap.set(cv.CAP_PROP_CONTRAST, 0.5) # 对比度
cap.set(cv.CAP_PROP_SHARPNESS, 25) # 锐度
cap.set(cv.CAP_PROP_AUTOFOCUS, 1) # 自动对焦
# 检查对焦状态
while True:
ret, frame = cap.read()
if not ret:
print("无法获取帧")
break
# 检查对焦是否完成
if check_focus(frame):
break
print("按下 's' 键拍照,点击鼠标左键保存图片,或按 'q' 键退出")
# 设置鼠标回调函数
cv.namedWindow('Press "s" to capture an image')
cv.setMouseCallback('Press "s" to capture an image', mouse_click)
while True:
ret, frame = cap.read()
if not ret:
print("无法获取帧")
break
# 显示视频流
cv.imshow('Press "s" to capture an image', frame)
key = cv.waitKey(1) & 0xFF
# 按下 's' 键保存图片
if key == ord('s'):
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
image_filename = os.path.join(image_dir, f"image_{timestamp}.png")
cv.imwrite(image_filename, frame) # 使用默认参数保存PNG,确保无损
print(f"图片已保存为 {image_filename}")
save_count += 1
print(f"已保存图片总数: {save_count}")
elif key == ord('q'): # 按下 'q' 键退出
print("退出程序")
break
# 释放资源
cap.release()
cv.destroyAllWindows()
def check_focus(frame):
# 使用边缘检测来检查对焦情况
gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
edges = cv.Canny(gray, 50, 150)
edge_count = np.sum(edges > 0)
# 如果边缘数量足够多,则认为对焦完成
return edge_count > 10000
if __name__ == "__main__":
capture_image()
一、RK3588环境准备
在这之前,需要先准备RK3588的环境,环境如下:
三个文件的链接放在这:ultralytics_yolov8,rknn_model_zoo ,rknn-toolkit2
在该路径下:
/home/sxj/YOLO/RKNN2_2.1.0/2.3.0/release/rknn-toolkit2-v2.3.0-2024-11-08/rknn-toolkit2/packages/arm64
1、找到这两个文件拷贝到RK3588上:
执行命令安装环境:
pip install -r arm64_requirements_cp310.txt
pip install rknn_toolkit2-2.3.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
如果是PC端:
pip install -r requirements_cp310-2.3.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install rknn_toolkit2-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
验证是否安装成功
在终端输入以下内容不报错代表成功
python
from rknn.api import RKNN
到此环境安装结束!
如果网速太慢,可以配置清华源进行下载,分别运行以下命令进行操作即可完成配置。(设置默认)
python -m pip install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --upgrade pip
pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
接下来模型转换
二、模型转换环境及生成onnx模型
1、查看模型
在训练得到PT模型后,我们可以用netron软件打开模型,看一下结构
激活函数已经为ReLU了
2. 下载瑞芯微修改后的源码,用于导出onnx:
git clone https://github.com/airockchip/ultralytics_yolov8.git
3.2
修改/home/sxj/yolov8-2/ultralytics_yolov8/ultralytics/cfg/default.yaml文件,改成自己模型名称
model: /home/sxj/yolov8-2/ultralytics_yolov8/best29.pt # (str, optional) path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: /home/sxj/yolov8-2/ultralytics_yolov8/ultralytics/cfg/datasets/coco.yaml # (str, optional) path to data file, i.e. coco8.yaml
3.3
同样修改ultralytics_yolov8/ultralytics/cfg/models/v8/yolov8.yaml
中的nc
为标签种类数
然后在ultralytics_yolov8主目录下执行:
conda activate yolo8
export PYTHONPATH=./
python ./ultralytics/engine/exporter.py
执行完毕后会在主目录下生成best29.onnx模型
3.4验证onnx模型(这个也可以放在后面因为需要安装rknn-toolkit2)
进入rknn230环境下
conda activate rknn230
在rknn_moodel_zoo/examples/yolov8/python/中执行以下代码测试onnx是否正确
修改yolov8.py中的CLASSES和coco_id_list
修改/rknn_model_zoo/examples/yolov8/model/coco_80_labels_list.txt为自己的标签种类
修改同目录下的dataset.txt为自己数据集中的一张照片
python yolov8.py --model_path best29.onnx --img_show
可以看到有正常输出检测框坐标
*******************************************************************************************************
rknn修改后的ultralytics_yolov8项目到本地:ultralytics_yolov8
ONNX转换为RKNN模型需要使用官方rknn_model_zoo工具:rknn_model_zoo-2.2.0
该处环境部署代码使用到官方rknn-toolkit2工具:rknn-toolkit2
(这里使用的都是最新版本)
板端代码使用优化修改后的开源代码:rknn3588-yolov8
报错:
转化过程中出现这种问题就是没有转换成功:
PyTorch: starting from 'best26.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 33, 8400) (6.0 MB)
RKNN: starting export with rknn-toolkit2...
ONNX: starting export with onnx 1.17.0 opset 19...
ONNX: slimming with onnxslim 0.1.48...
解决方法:
1、下载修改后的rknn修改后的ultralytics_yolov8项目到本地,在该目录下转换,这个是瑞星微修改后的
运行之后生成有如下就成功了
Ultralytics YOLOv8.2.82 🚀 Python-3.10.16 torch-2.4.0+cu121 CPU (13th Gen Intel Core(TM) i7-13620H)
Model summary (fused): 168 layers, 3,011,303 parameters, 0 gradients, 8.1 GFLOPs
PyTorch: starting from '/home/sxj/yolov8-2/ultralytics_yolov8/best29.pt' with input shape (16, 3, 640, 640) BCHW and output shape(s) ((16, 64, 80, 80), (16, 29, 80, 80), (16, 1, 80, 80), (16, 64, 40, 40), (16, 29, 40, 40), (16, 1, 40, 40), (16, 64, 20, 20), (16, 29, 20, 20), (16, 1, 20, 20)) (6.0 MB)
RKNN: starting export with torch 2.4.0+cu121...
RKNN: feed /home/sxj/yolov8-2/ultralytics_yolov8/best29.onnx to RKNN-Toolkit or RKNN-Toolkit2 to generate RKNN model.
转换成功的环境如下:
(yolo8) sxj@sxj:~/yolov8-2/ultralytics_yolov8$ pip list
Package Version
------------------------ -----------
aiofiles 23.2.1
annotated-types 0.7.0
anyio 4.8.0
certifi 2024.12.14
charset-normalizer 3.4.1
click 8.1.8
coloredlogs 15.0.1
contourpy 1.3.1
cycler 0.12.1
exceptiongroup 1.2.2
fast-histogram 0.14
fastapi 0.115.8
ffmpy 0.5.0
filelock 3.17.0
flatbuffers 25.1.24
fonttools 4.55.5
fsspec 2024.12.0
gradio 5.15.0
gradio_client 1.7.0
h11 0.14.0
httpcore 1.0.7
httpx 0.28.1
huggingface-hub 0.28.1
humanfriendly 10.0
idna 3.10
Jinja2 3.1.5
kiwisolver 1.4.8
markdown-it-py 3.0.0
MarkupSafe 2.1.5
matplotlib 3.10.0
mdurl 0.1.2
mpmath 1.3.0
netron 8.1.3
networkx 3.4.2
numpy 1.26.4
nvidia-cublas-cu12 12.1.3.1
nvidia-cuda-cupti-cu12 12.1.105
nvidia-cuda-nvrtc-cu12 12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.0.2.54
nvidia-curand-cu12 10.3.2.106
nvidia-cusolver-cu12 11.4.5.107
nvidia-cusparse-cu12 12.1.0.106
nvidia-cusparselt-cu12 0.6.2
nvidia-nccl-cu12 2.20.5
nvidia-nvjitlink-cu12 12.4.127
nvidia-nvtx-cu12 12.1.105
onnx 1.17.0
onnxruntime 1.20.1
onnxruntime-gpu 1.20.1
onnxslim 0.1.48
opencv-python 4.11.0.86
orjson 3.10.15
packaging 24.2
pandas 2.2.3
pillow 11.1.0
pip 24.2
protobuf 4.25.4
psutil 6.1.1
py-cpuinfo 9.0.0
pydantic 2.10.6
pydantic_core 2.27.2
pydub 0.25.1
Pygments 2.19.1
pyparsing 3.2.1
python-dateutil 2.9.0.post0
python-multipart 0.0.20
pytz 2024.2
PyYAML 6.0.2
requests 2.32.3
rich 13.9.4
rknn-toolkit2 2.3.0
ruamel.yaml 0.18.10
ruamel.yaml.clib 0.2.12
ruff 0.9.5
safehttpx 0.1.6
scipy 1.15.1
seaborn 0.13.2
semantic-version 2.10.0
setuptools 75.1.0
shellingham 1.5.4
six 1.17.0
sniffio 1.3.1
starlette 0.45.3
sympy 1.13.1
tomlkit 0.13.2
torch 2.4.0
torchvision 0.19.0
tqdm 4.67.1
triton 3.0.0
typer 0.15.1
typing_extensions 4.12.2
tzdata 2025.1
ultralytics 8.3.72
ultralytics-thop 2.0.14
urllib3 2.3.0
uvicorn 0.34.0
websockets 14.2
wheel 0.44.0
如果报如下
RKNN: starting export with rknn-toolkit2...
requirements: Ultralytics requirement ['rknn-toolkit2'] not found, attempting AutoUpdate...
说明没有转成功
3、ONNX转RKNN环境
这里我们先创建环境:(这里是在PC端的环境)
conda create -n rknn230 python=3.10
conda activate rknn230
现在需要用到rknn-toolkit2-2.3.0文件
进入rknn-toolkit2-2.3.0\rknn-toolkit2-2.3.0\rknn-toolkit2\packages文件夹下
我的目录是:
/home/sxj/YOLO/RKNN2_2.1.0/2.3.0/release/rknn-toolkit2-v2.3.0-2024-11-08/rknn-toolkit2/packages/x86_64
终端输入:
pip install -r requirements_cp310-2.3.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
然后再输入:
pip install rknn_toolkit2-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
这里我们的转rknn环境就配置完成了!!!
三、转换RKNN模型详细流程
1、修改yolov8.py
先进入rknn_model_zoo-2.1.0\examples\yolov8\python文件夹,先打开yolov8.py,进行适配参数修改:
2、然后修改convert.py
如下所示:
修改完成后,将我们之前得到的onnx模型复制到python文件夹下:
打开终端,激活rknn230环境,输入命令:
python convert.py best13.onnx rk3588
非量化版本
python convert.py best71.onnx rk3588 fp
结果如下:
转换后的rknn模型已保存在python文件夹下
用netron打开我们的rknn模型,看一下结构:
到此模型转换完成!!!
参考:【YOLOv8部署至RK3588】模型训练→转换RKNN→开发板部署_yolov8转rknn-CSDN博客
四、模型部署(开发板RK3588上)
首先进入/rknn_model_zoo/examples/yolov8/python下:
修改yolov8.py里面改成自己的
CLASSES = ('sheqian','yiyongjiaodai','yiyongmianqian')
coco_id_list = [1, 2, 3]
1、python运行
cd python
使用onnx模型推理:
python yolov8.py --model_path ../model/yolov8n.onnx --img_show
python yolov8.py --model_path ../model/best13.onnx --img_show
使用rknn模型推理:(常用)
python yolov81.py --model_path best30.rknn --target rk3588 --img_show
python yolov8.py --model_path ../model/yolov8n.rknn --target rk3588 --img_show
2、报错内容:
python yolov8.py --model_path ../model/yolov8n.rknn --target rk3588 --img_show
I rknn-toolkit2 version: 2.3.0
--> Init runtime environment
I target set by user is: rk3588
E RKNN: [11:28:24.309] 6, 1
E RKNN: [11:28:24.309] Invalid RKNN model version 6
E RKNN: [11:28:24.309] rknn_init, load model failed!
E init_runtime: Traceback (most recent call last):
File "rknn/api/rknn_log.py", line 344, in rknn.api.rknn_log.error_catch_decorator.error_catch_wrapper
File "rknn/api/rknn_base.py", line 2483, in rknn.api.rknn_base.RKNNBase.init_runtime
File "rknn/api/rknn_runtime.py", line 427, in rknn.api.rknn_runtime.RKNNRuntime.build_graph
Exception: RKNN init failed. error code: RKNN_ERR_FAIL
3、解决方法:
解决方案
下载rknpu2
git clone https://github.com/rockchip-linux/rknpu2
将下面的so文件复制到/usr/lib/下
sudo cp rknpu2/runtime/RK3588/Linux/librknn_api/aarch64/librknnrt.so /usr/lib/librknnrt.so
或
sudo cp librknnrt.so /usr/lib/librknnrt.so
再次运行:
python yolov81.py --model_path best30.rknn --target rk3588 --img_show
运行成功!!!
五、实时视频流
暂时只能PC端跑
from ultralytics import YOLO
import cv2
from cv2 import getTickCount, getTickFrequency
yolo=YOLO('./best26.pt')
#摄像头实时检测
cap = cv2.VideoCapture(2)
while cap.isOpened():
loop_start = getTickCount() #记录循环开始的时间,用于计算每一帧的处理时间
success, frame = cap.read() # 读取摄像头的一帧图像
if success:
results = yolo.predict(source=frame) # 对当前帧进行目标检测并显示结果
annotated_frame = results[0].plot() #将检测结果绘制在图像上,得到带有目标框的图像。
# 显示程序
loop_time = getTickCount() - loop_start #计算处理一帧图像所花费的时间。
total_time = loop_time / (getTickFrequency()) #将处理时间转换为秒数。
FPS = int(1 / total_time) #FPS计算
# 在图像左上角添加FPS文本
fps_text = f"FPS: {FPS:.2f}"#构造显示帧率的文本字符串
font = cv2.FONT_HERSHEY_SIMPLEX#选择字体类型
font_scale = 1 #设置字体的缩放比例
font_thickness = 2 #设置字体的粗细
text_color = (0, 0, 255) # 红色
text_position = (10, 30) # 左上角位置
cv2.putText(annotated_frame, fps_text, text_position, font, font_scale, text_color, font_thickness)
cv2.imshow('Real-time detection', annotated_frame)
# 通过按下 'q' 键退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release() # 释放摄像头资源
cv2.destroyAllWindows() # 关闭OpenCV窗口
python video.py
import cv2
import time
import numpy as np
from queue import Queue
from concurrent.futures import ThreadPoolExecutor
from rknnlite.api import RKNNLite
# 常量定义
OBJ_THRESH, NMS_THRESH, IMG_SIZE = 0.45, 0.45, 640
CLASSES = (
'rusuan'
)
# ====================== 从rknnpool.py合并的代码 ======================
def initRKNN(rknnModel="./rknnModel/yolov5s.rknn", id=0):
rknn_lite = RKNNLite()
ret = rknn_lite.load_rknn(rknnModel)
if ret != 0:
print("Load RKNN rknnModel failed")
exit(ret)
if id == 0:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
elif id == 1:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_1)
elif id == 2:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_2)
elif id == -1:
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
else:
ret = rknn_lite.init_runtime()
if ret != 0:
print("Init runtime environment failed")
exit(ret)
print(rknnModel, "\t\tdone")
return rknn_lite
def initRKNNs(rknnModel="./rknnModel/yolov5s.rknn", TPEs=1):
rknn_list = []
for i in range(TPEs):
rknn_list.append(initRKNN(rknnModel, i % 3))
return rknn_list
class rknnPoolExecutor():
def __init__(self, rknnModel, TPEs, func):
self.TPEs = TPEs
self.queue = Queue()
self.rknnPool = initRKNNs(rknnModel, TPEs)
self.pool = ThreadPoolExecutor(max_workers=TPEs)
self.func = func
self.num = 0
def put(self, frame):
self.queue.put(self.pool.submit(
self.func, self.rknnPool[self.num % self.TPEs], frame))
self.num += 1
def get(self):
if self.queue.empty():
return None, False
fut = self.queue.get()
return fut.result(), True
def release(self):
self.pool.shutdown()
for rknn_lite in self.rknnPool:
rknn_lite.release()
# ====================== 原有video.py代码 ======================
def filter_boxes(boxes, box_confidences, box_class_probs):
box_confidences = box_confidences.reshape(-1)
candidate, class_num = box_class_probs.shape
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 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]
return np.array(keep)
def dfl(position):
n, c, h, w = position.shape
p_num = 4
mc = c // p_num
y = position.reshape(n, p_num, mc, h, w)
e_y = np.exp(y - np.max(y, axis=2, keepdims=True))
y = e_y / np.sum(e_y, axis=2, keepdims=True)
acc_metrix = np.arange(mc).reshape(1, 1, mc, 1, 1)
y = (y * acc_metrix).sum(2)
return y
def box_process(position):
grid_h, grid_w = position.shape[2:4]
col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
col = col.reshape(1, 1, grid_h, grid_w)
row = row.reshape(1, 1, grid_h, grid_w)
grid = np.concatenate((col, row), axis=1)
stride = np.array([IMG_SIZE//grid_h, IMG_SIZE//grid_w]).reshape(1,2,1,1)
position = dfl(position)
box_xy = grid + 0.5 - position[:,0:2,:,:]
box_xy2 = grid + 0.5 + position[:,2:4,:,:]
xyxy = np.concatenate((box_xy*stride, box_xy2*stride), axis=1)
return xyxy
def yolov8_post_process(input_data):
boxes, scores, classes_conf = [], [], []
defualt_branch = 3
pair_per_branch = len(input_data) // defualt_branch
for i in range(defualt_branch):
boxes.append(box_process(input_data[pair_per_branch*i]))
classes_conf.append(input_data[pair_per_branch*i+1])
scores.append(np.ones_like(input_data[pair_per_branch*i+1][:,:1,:,:], dtype=np.float32))
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]
classes_conf = [sp_flatten(_v) for _v in classes_conf]
scores = [sp_flatten(_v) for _v in scores]
boxes = np.concatenate(boxes)
classes_conf = np.concatenate(classes_conf)
scores = np.concatenate(scores)
boxes, classes, scores = filter_boxes(boxes, scores, classes_conf)
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
return np.concatenate(nboxes), np.concatenate(nclasses), np.concatenate(nscores)
def draw(image, boxes, scores, classes, ratio, padding):
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
top = int((top - padding[0]) / ratio[0])
left = int((left - padding[1]) / ratio[1])
right = int((right - padding[0]) / ratio[0])
bottom = int((bottom - padding[1]) / ratio[1])
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
cv2.putText(image, f'{CLASSES[cl]} {score:.2f}',
(top, left - 6),
cv2.FONT_HERSHEY_SIMPLEX,
0.6, (0, 0, 255), 2)
def letterbox(im, new_shape=(640, 640), color=(255, 255, 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])
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, (r, r), (left, top)
def myFunc(rknn_lite, IMG):
IMG2 = cv2.cvtColor(IMG, cv2.COLOR_BGR2RGB)
IMG2, ratio, padding = letterbox(IMG2)
IMG2 = np.expand_dims(IMG2, 0)
outputs = rknn_lite.inference(inputs=[IMG2], data_format=['nhwc'])
boxes, classes, scores = yolov8_post_process(outputs)
if boxes is not None:
draw(IMG, boxes, scores, classes, ratio, padding)
return IMG
if __name__ == "__main__":
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
modelPath = "./best31.rknn" # 修改为实际模型路径
TPEs = 6
pool = rknnPoolExecutor(
rknnModel=modelPath,
TPEs=TPEs,
func=myFunc)
if cap.isOpened():
for _ in range(TPEs + 1):
ret, frame = cap.read()
if not ret:
cap.release()
del pool
exit(-1)
pool.put(frame)
frames, loopTime, initTime = 0, time.time(), time.time()
while cap.isOpened():
frames += 1
ret, frame = cap.read()
if not ret:
break
pool.put(frame)
frame, flag = pool.get()
if not flag:
break
cv2.imshow('YOLOv8 Detection', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if frames % 30 == 0:
print(f"30帧平均帧率: {30 / (time.time() - loopTime):.2f} 帧")
loopTime = time.time()
print(f"总平均帧率: {frames / (time.time() - initTime):.2f} 帧")
cap.release()
cv2.destroyAllWindows()
pool.release()
python3 video2.py
import cv2
import time
import numpy as np
from rknnpool import rknnPoolExecutor
# 常量定义
OBJ_THRESH, NMS_THRESH, IMG_SIZE = 0.45, 0.45, 640
CLASSES = (
'tongkongbi'
)
def filter_boxes(boxes, box_confidences, box_class_probs):
"""Filter boxes with object threshold."""
box_confidences = box_confidences.reshape(-1)
candidate, class_num = box_class_probs.shape
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 nms_boxes(boxes, scores):
"""Suppress non-maximal 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]
return np.array(keep)
def dfl(position):
"""Distribution Focal Loss (DFL)."""
n, c, h, w = position.shape
p_num = 4
mc = c // p_num
y = position.reshape(n, p_num, mc, h, w)
# Vectorized softmax
e_y = np.exp(y - np.max(y, axis=2, keepdims=True))
y = e_y / np.sum(e_y, axis=2, keepdims=True)
acc_metrix = np.arange(mc).reshape(1, 1, mc, 1, 1)
y = (y * acc_metrix).sum(2)
return y
def box_process(position):
"""Process box coordinates."""
grid_h, grid_w = position.shape[2:4]
col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
col = col.reshape(1, 1, grid_h, grid_w)
row = row.reshape(1, 1, grid_h, grid_w)
grid = np.concatenate((col, row), axis=1)
stride = np.array([IMG_SIZE//grid_h, IMG_SIZE//grid_w]).reshape(1,2,1,1)
position = dfl(position)
box_xy = grid + 0.5 - position[:,0:2,:,:]
box_xy2 = grid + 0.5 + position[:,2:4,:,:]
xyxy = np.concatenate((box_xy*stride, box_xy2*stride), axis=1)
return xyxy
def yolov8_post_process(input_data):
"""YOLOv8 post-processing."""
boxes, scores, classes_conf = [], [], []
defualt_branch = 3
pair_per_branch = len(input_data) // defualt_branch
for i in range(defualt_branch):
boxes.append(box_process(input_data[pair_per_branch*i]))
classes_conf.append(input_data[pair_per_branch*i+1])
scores.append(np.ones_like(input_data[pair_per_branch*i+1][:,:1,:,:], dtype=np.float32))
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]
classes_conf = [sp_flatten(_v) for _v in classes_conf]
scores = [sp_flatten(_v) for _v in scores]
boxes = np.concatenate(boxes)
classes_conf = np.concatenate(classes_conf)
scores = np.concatenate(scores)
boxes, classes, scores = filter_boxes(boxes, scores, classes_conf)
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
return np.concatenate(nboxes), np.concatenate(nclasses), np.concatenate(nscores)
def draw(image, boxes, scores, classes, ratio, padding):
"""Draw bounding boxes on image."""
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
top = int((top - padding[0]) / ratio[0])
left = int((left - padding[1]) / ratio[1])
right = int((right - padding[0]) / ratio[0])
bottom = int((bottom - padding[1]) / ratio[1])
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
cv2.putText(image, f'{CLASSES[cl]} {score:.2f}',
(top, left - 6),
cv2.FONT_HERSHEY_SIMPLEX,
0.6, (0, 0, 255), 2)
def letterbox(im, new_shape=(640, 640), color=(255, 255, 0)):
"""Resize image with unchanged aspect ratio using padding."""
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])
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, (r, r), (left, top)
def myFunc(rknn_lite, IMG):
"""Image processing function for RKNN."""
IMG2 = cv2.cvtColor(IMG, cv2.COLOR_BGR2RGB)
IMG2, ratio, padding = letterbox(IMG2)
IMG2 = np.expand_dims(IMG2, 0)
outputs = rknn_lite.inference(inputs=[IMG2], data_format=['nhwc'])
boxes, classes, scores = yolov8_post_process(outputs)
detected_classes = []
if boxes is not None:
draw(IMG, boxes, scores, classes, ratio, padding)
# 收集检测到的类别
detected_classes = [CLASSES[cl] for cl in classes]
return IMG, detected_classes # 修改返回值
if __name__ == "__main__":
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
modelPath = "./best31.rknn" # 根据实际路径修改
TPEs = 6
pool = rknnPoolExecutor(
rknnModel=modelPath,
TPEs=TPEs,
func=myFunc)
# 初始化缓冲帧
if cap.isOpened():
for _ in range(TPEs + 1):
ret, frame = cap.read()
if not ret:
cap.release()
del pool
exit(-1)
pool.put(frame)
frames, loopTime, initTime = 0, time.time(), time.time()
while cap.isOpened():
frames += 1
ret, frame = cap.read()
if not ret:
break
pool.put(frame)
result, flag = pool.get() # 修改获取方式
if not flag:
break
# 解包结果
processed_frame, detected_classes = result
# 实时打印检测结果
if detected_classes:
print(f"帧 {frames} 检测到类别:")
for cls in set(detected_classes): # 去重后打印
count = detected_classes.count(cls)
print(f" {cls}: {count}个")
print("-" * 30)
cv2.imshow('YOLOv8 Detection', processed_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if frames % 30 == 0:
print(f"30帧平均帧率: {30 / (time.time() - loopTime):.2f} 帧")
loopTime = time.time()
print(f"总平均帧率: {frames / (time.time() - initTime):.2f} 帧")
cap.release()
cv2.destroyAllWindows()
pool.release()
终端输出:
参考:
【保姆级教程】从Yolov8训练模型到转化Onnx再转换为Rknn以及板端部署全记录_yolov8 rknn-CSDN博客