现在github下载
def config(self,
batch_size=100,
caffe_mean_file=None,
dtype='float32',
epochs=-1,
force_gray=None,
input_fitting='scale',
input_normalization=None,
mean_file=None,
model_data_format='zone',
model_quantize=None,
optimize='Default',
quantized_dtype='asymmetric_affine-u8',
quantized_moving_alpha=0.01,
quantized_algorithm='normal',
quantized_divergence_nbins=1024,
mmse_epoch=3,
random_brightness=None,
random_contrast=None,
random_crop=None,
random_flip=None,
random_mirror=None,
reorder_channel=None,
restart=False,
samples=-1,
need_horizontal_merge=False,
deconv_merge=True,
conv_mul_merge=True,
quantized_hybrid=False,
output_optimize=0,
remove_tensorflow_output_permute=False,
optimization_level=3,
target_platform=None,
mean_values=None,
std_values=None,
channel_mean_value=None,
force_builtin_perm=False,
do_sparse_network=True):
rknn.build(do_quantization,dataset,pre_compile,rknn_batch_size)
pre_compile:预编译开关,如果设置成 True,可以减小模型大小,及模型在硬件设备
上的首次启动速度。但是打开这个开关后,构建出来的模型就只能在硬件平台上运
行,无法通过模拟器进行推理或性能评估。如果硬件有更新,则对应的模型要重新构
建。
rknn.load_rknn(path,load_model_in_npu)
是否直接加载 npu 中的 rknn 模型。其中 path 为 rknn 模型在 npu
35
中的路径。只有当 RKNN-Toolkit 运行在 RK3399Pro Linux 开发板或连有 NPU 设备
的 PC 上时才可以设为 True。默认值为 False。
rknn.init_runtime(target,device_id,perf_debug,eval_mem,async_mode)
target:目标硬件平台,目前支持“rk3399pro”、“rk1806”、“rk1808”、“rv1109”、
“rv1126”。默认为 None,即在 PC 使用工具时,模型在模拟器上运行,在 RK3399Pro
Linux 开发板运行时,模型在 RK3399Pro 自带 NPU 上运行,否则在设定的 target 上
运行。其中“rk1808”包含了 TB-RK1808 AI 计算棒
perf_debug:进行性能评估时是否开启 debug 模式。在 debug 模式下,可以获取到每
一层的运行时间,否则只能获取模型运行的总时间。默认值为 False。
eval_mem: 是否进入内存评估模式。进入内存评估模式后,可以调用 eval_memory 接
口获取模型运行时的内存使用情况。默认值为 False。
async_mode:是否使用异步模式。调用推理接口时,涉及设置输入图片、模型推理、
36
获取推理结果三个阶段。如果开启了异步模式,设置当前帧的输入将与推理上一帧同
时进行,所以除第一帧外,之后的每一帧都可以隐藏设置输入的时间,从而提升性能。
在异步模式下,每次返回的推理结果都是上一帧的。该参数的默认值为 False。
rknn.inference(inputs,data_type,data_format,inputs_pass_through)
inputs:待推理的输入,如经过 cv2 处理的图片。格式是 ndarray list。
data_type:输入数据的类型,可填以下值: ’float32’, ‘float16’, ‘int8’, ‘uint8’, ‘int16’。默认值为’uint8’。
data_format:数据模式,可以填以下值: “nchw”, “nhwc”。默认值为’nhwc’。这两个的不同之处在于 channel 放置的位置。
inputs_pass_through: 将输入透传给 NPU 驱动。非透传模式下,在将输入传给 NPU 驱动之前,工具会对输入进行减均值、除方差等操作;而透传模式下,不会做这些操作。
这个参数的值是一个数组,比如要透传 input0,不透彻 input1,则这个参数的值为[1,0]。默认值为 None,即对所有输入都不透传。
import os
import numpy as np
import cv2
from rknn.api import RKNN
ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'yolov.rknn'
IMG_PATH = '3.jpg'
DATASET = 'lh.txt'
NMS_THRESH=0.5
IMG_SIZE=640
strides=[8,16,32]
anchors=np.array([10,13, 16,30, 33,23,30,61, 62,45, 59,119,116,90, 156,198, 373,326])
ahchors=anchors.reshape((3,1,3,1,1,2))
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]
keep = np.array(keep)
return keep
def sigmoid(x):
s = 1 / (1 + np.exp(-x))
return s
def make_grid(nx=20, ny=20):
X, Y = np.meshgrid(np.arange(0, nx, 1),np.arange(0, ny, 1))
return np.stack((X, Y), 2).reshape((1, 1, ny, nx, 2))
def my_process(outputs,i):
b,_,ny,nx=outputs.shape
output = np.transpose(outputs.reshape(1, 3, 7, ny, nx),(0, 1, 3, 4, 2))
output=np.ascontiguousarray(output)
y = sigmoid(output)
grid = make_grid(nx, ny)
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid) * strides[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchors[i]
y=y.reshape(-1,7)
return y
def letterbox(im, new_shape=(640, 640), color=(0, 0, 0)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
ratio = r, r # width, height ratios
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] # wh padding
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
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) # add border
return im, ratio, (dw, dh)
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN()
if not os.path.exists(ONNX_MODEL):
print('model not exist')
exit(-1)
# pre-process config
print('--> Config model')
rknn.config(reorder_channel='0 1 2',
mean_values=[[0, 0, 0]],
std_values=[[255, 255, 255]],
optimization_level=3)
print('done')
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
print('Load yolov5 failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False, dataset=DATASET)
if ret != 0:
print('Build yolov5 failed!')
exit(ret)
print('done')
# Export RKNN model
print('--> Export RKNN model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
print('Export yolov5rknn failed!')
exit(ret)
print('done')
# init runtime environment
print('--> Init runtime environment')
# ret = rknn.init_runtime()
ret = rknn.init_runtime(device_id="0")
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# Set inputs
img = cv2.imread(IMG_PATH)
img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
img = cv2.resize(img, (640, 640), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
color=[127,210,155]
# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img])
results=[]
for i in range(3):
y=my_process(outputs[i],i)
results.append(y)
result=np.concatenate((results[0],results[1]),axis=0)
result=np.concatenate((result,results[2]),axis=0)
for i in range(result.shape[0]):
if result[i][4]>0.3:
xmin = int(result[i][0]-result[i][2]/2)
ymin = int(result[i][1]-result[i][3]/2)
xmax = int(result[i][0]+result[i][2]/2)
ymax = int(result[i][1]+result[i][3]/2)
print(xmin,ymin,xmax,ymax)
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, thickness=3, lineType=cv2.LINE_AA)
cv2.imwrite("11.jpg", img)
rknn.release()
要注意的是没有做nms操作,并且resize图片使用的是letterbox,但最后没有按此方式映射坐标,自己写个简单影射即可,即坐标%ratio-dw即可。主要把torch的后处理过程用python实现一遍详见my_process过程。
import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
ONNX_MODEL = 'best1.onnx'
RKNN_MODEL = 'yolov1.rknn'
IMG_PATH = '6.jpg'
DATASET = 'lh.txt'
OBJ_THRESH = 0.5
NMS_THRESH = 0.6
IMG_SIZE = 640
def nms_boxes(boxes, scores):
"""Suppress non-maximal boxes.
# Arguments
boxes: ndarray, boxes of objects.
scores: ndarray, scores of objects.
# Returns
keep: ndarray, index of effective 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]
keep = np.array(keep)
return keep
def plt_img(outputs,img):
output=outputs[0,:,:]
color=[127,0,255]
for i in range(output.shape[0]):
if output[i][4]>0.3:
print("ok")
xmin=int(output[i][0]-output[i][2]/2)
ymin = int(output[i][1]-output[i][3]/2)
xmax = int(output[i][0]+output[i][2]/2)
ymax = int(output[i][1]+output[i][3]/2)
print(xmin,ymin,xmax,ymax)
cv2.rectangle(img, (xmin,ymin), (xmax,ymax), color, thickness=3, lineType=cv2.LINE_AA)
cv2.imwrite("1.jpg",img)
def letterbox(im, new_shape=(640, 640), color=(0, 0, 0)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
ratio = r, r # width, height ratios
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] # wh padding
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
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) # add border
return im, ratio, (dw, dh)
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN()
if not os.path.exists(ONNX_MODEL):
print('model not exist')
exit(-1)
# pre-process config
print('--> Config model')
rknn.config(reorder_channel='0 1 2',
mean_values=[[0, 0, 0]],
std_values=[[255, 255, 255]],
optimization_level=3)
print('done')
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
print('Load yolov5 failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False, dataset=DATASET)
if ret != 0:
print('Build yolov5 failed!')
exit(ret)
print('done')
# Export RKNN model
print('--> Export RKNN model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
print('Export yolov5rknn failed!')
exit(ret)
print('done')
# init runtime environment
print('--> Init runtime environment')
# ret = rknn.init_runtime()
ret = rknn.init_runtime(device_id="0")
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# Set inputs
img = cv2.imread(IMG_PATH)
img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img])
plt_img(outputs[0], img)
print("Ok")
rknn.release()
只有一个输出的话就没有后处理过程直接将xywh2xyxy即可,不要忘记nms和坐标影射。