59 lines
1.5 KiB
Python
59 lines
1.5 KiB
Python
from rknn.api import RKNN
|
||
import cv2
|
||
import numpy as np
|
||
|
||
|
||
# 初始化 RKNN
|
||
rknn = RKNN()
|
||
|
||
# 配置参数（关键！）
|
||
rknn.config(
|
||
target_platform="rk3588", # 根据实际芯片型号修改
|
||
mean_values=[[0, 0, 0]], # YOLOv8 输入为 0-255，无需归一化
|
||
std_values=[[255, 255, 255]], # 输入数据除以 255（即 0-1 范围）
|
||
# quant_img_RGB2BGR=True,
|
||
optimization_level=3, # 最高优化级别
|
||
)
|
||
|
||
# 加载 ONNX
|
||
ret = rknn.load_onnx(model="/home/admin-root/haotian/rk3588/pytorch模型转rknn/models/yolov8m.onnx")
|
||
assert ret == 0, "加载 ONNX 失败！"
|
||
|
||
# 转换模型
|
||
ret = rknn.build(
|
||
do_quantization=False, # 启用量化
|
||
# dataset="dataset.txt", # 校准数据路径
|
||
)
|
||
assert ret == 0, "转换 RKNN 失败！"
|
||
|
||
# 导出 RKNN
|
||
ret = rknn.export_rknn("/home/admin-root/haotian/rk3588/pytorch模型转rknn/models/yolov8m.rknn")
|
||
assert ret == 0, "导出 RKNN 失败！"
|
||
|
||
|
||
# Set inputs
|
||
img = cv2.imread('/home/admin-root/haotian/rk3588/pytorch模型转rknn/images/bus.jpg')
|
||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||
img.resize((3, 640, 640))
|
||
img = np.expand_dims(img, 0)
|
||
|
||
# Init runtime environment
|
||
print('--> Init runtime environment')
|
||
ret = rknn.init_runtime()
|
||
if ret != 0:
|
||
print('Init runtime environment failed!')
|
||
exit(ret)
|
||
print('done')
|
||
|
||
# Inference
|
||
print('--> Running model')
|
||
outputs = rknn.inference(inputs=[img], data_format=['nchw'])
|
||
np.save('./tflite_mobilenet_v1_0.npy', outputs[0])
|
||
print(len(outputs))
|
||
print('done')
|
||
|
||
rknn.release()
|
||
|
||
|
||
|