模型推理¶
本章节主要讲解如何在Toybrick rk3399Pro开发板上调用RKNN Python API进行模型推理。
API调用流程
示例
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | import numpy as np
from PIL import Image
from rknn.api import RKNN
# 解析模型的输出,获得概率最大的手势和对应的概率
def get_predict(probability):
data = probability[0][0]
data = data.tolist()
max_prob = max(data)
return data.index(max_prob), max_prob;
def load_model():
# 创建RKNN对象
rknn = RKNN()
# 载入RKNN模型
print('-->loading model')
rknn.load_rknn('./digital_gesture.rknn')
print('loading model done')
# 初始化RKNN运行环境
print('--> Init runtime environment')
ret = rknn.init_runtime(target='rk3399pro')
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
return rknn
def predict(rknn):
im = Image.open("../picture/6_7.jpg") # 加载图片
im = im.resize((64, 64),Image.ANTIALIAS) # 图像缩放到64x64
mat = np.asarray(im.convert('RGB')) # 转换成RGB格式
outputs = rknn.inference(inputs=[mat]) # 运行推理,得到推理结果
pred, prob = get_predict(outputs) # 将推理结果转化为可视信息
print(prob)
print(pred)
if __name__=="__main__":
rknn = load_model()
predict(rknn)
rknn.release()
|