62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
from rknn.api import RKNN
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
# 初始化 RKNN
|
|
rknn = RKNN()
|
|
|
|
# 配置参数（关键！）
|
|
rknn.config(
|
|
target_platform="rk3588", # 根据实际芯片型号修改
|
|
# mean_values=[[123.675, 116.28, 103.53]], # 默认输入是255, 使用这些值归一化到0-1
|
|
# std_values=[[58.395, 57.12, 57.375]],
|
|
mean_values=[[0, 0, 0]], # 默认输入是255, 使用这些值归一化到0-1
|
|
std_values=[[255, 255, 255]],
|
|
quant_img_RGB2BGR=True,
|
|
optimization_level=3, # 最高优化级别
|
|
)
|
|
|
|
# 加载 ONNX, 只支持固定输入的onnx模型
|
|
ret = rknn.load_onnx(model="/home/admin-root/haotian/康达瑞贝斯机器狗/rec_shape_20250815.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/康达瑞贝斯机器狗/rec_shape_20250815_0m.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.resize((3, 48, 320))
|
|
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()
|
|
|
|
|
|
|