- if __name__ == '__main__':
- # Create RKNN object
- rknn = RKNN(verbose=True)
- # pre-process config
- print('--> Config model')
- rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform='rk3568', optimization_level=0)
- print('done')
- # Load ONNX model
- print('--> Loading model')
- ret = rknn.load_onnx(model=ONNX_MODEL)
- if ret != 0:
- print('Load model failed!')
- exit(ret)
- print('done')
- # Build model
- print('--> Building model')
- # ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
- ret = rknn.build(do_quantization=False, dataset=DATASET)
- if ret != 0:
- print('Build model failed!')
- exit(ret)
- print('done')
- # Accuracy test 连扳推理测试在板子上的查看每一层的推理精度是否下降
- # ret = rknn.accuracy_**ysis(
- # inputs=[IMG_PATH],
- # output_dir="snapshat",
- # target="rk3568")
-
- if ret != 0:
- print('Accuracy test failed!')
- exit(ret)
- print('done')
- # Export RKNN model
- print('--> Export rknn model')
- ret = rknn.export_rknn(RKNN_MODEL)
- if ret != 0:
- print('Export rknn model failed!')
- exit(ret)
- print('done')
- # Init runtime environment
- print('--> Init runtime environment')
- ret = rknn.init_runtime()
- 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)
- img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
- # Inference
- print('--> Running model')
- img2 = np.expand_dims(img, 0)
- print("++++++++++++++++")
- print(img2.shape)
- print("++++++++++++++++")
- outputs = rknn.inference(inputs=[img2], data_format=['nhwc'])
- # np.save('./onnx_yolov5_0.npy', outputs[0])
- # np.save('./onnx_yolov5_1.npy', outputs[1])
- # np.save('./onnx_yolov5_2.npy', outputs[2])
-
- print('done')
- # post process
- input0_data = outputs[0]
- input1_data = outputs[1]
- input2_data = outputs[2]
- input0_data = input0_data.reshape([3, -1]+list(input0_data.shape[-2:]))
- input1_data = input1_data.reshape([3, -1]+list(input1_data.shape[-2:]))
- input2_data = input2_data.reshape([3, -1]+list(input2_data.shape[-2:]))
- input_data = list()
- input_data.append(np.transpose(input0_data, (2, 3, 0, 1)))
- input_data.append(np.transpose(input1_data, (2, 3, 0, 1)))
- input_data.append(np.transpose(input2_data, (2, 3, 0, 1)))
- boxes, classes, scores = yolov5_post_process(input_data)
- img_1 = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
- if boxes is not None:
- draw(img_1, boxes, scores, classes)
- cv2.imwrite('result.jpg', img_1)
- print('Save results to result.jpg!')
- rknn.release() 这是视频数据加载过程
复制代码
欢迎光临 Toybrick (https://t.rock-chips.com/) | Powered by Discuz! X3.3 |