使用rknn-toolkit2把YOLOV5部署到OK3588上
- 虚拟环境搭建
- 软件包安装
- 在PC机上运行yolov5目标检测
虚拟环境搭建
首先在PC的ubuntu系统安装虚拟环境:
我的服务器是ubuntu18.04版本,所以安装python3.6
conda create -n ok3588 python=3.6
需要键盘输入y,然后完成虚拟环境安装。
其他系统的对应关系:
Ubuntu 18.04 python 3.6 / Ubuntu 20.04 python 3.8 / Ubuntu 22.04 python 3.10
软件包安装
进入虚拟环境
conda activate ok3588
首先安装正确的pip版本
curl https://bootstrap.pypa.io/pip/3.6/get-pip.py | python -
下载github 项目
git clone https://github.com/rockchip-linux/rknn-toolkit2
cd rknn-toolkit2/doc
pip install -r requirements_cp36-1.5.0.txt -i https://mirror.baidu.com/pypi/simple
cd rknn-toolkit2/packages
pip install rknn_toolkit2-1.5.0+1fa95b5c-cp36-cp36m-linux_x86_64.whl
在PC机上运行yolov5目标检测
cd rknn-toolkit2/examples/onnx/yolov5
python test.py
截图如下
推理前的图片:
推理后加上box的图片:
推理代码和注释:
if __name__ == '__main__':# Create RKNN objectrknn = RKNN(verbose=True)# pre-process config 配置数据参数print('--> Config model')rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]])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)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN model 导出RKNN模型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 初始化runtime环境print('--> Init runtime environment')ret = rknn.init_runtime()# ret = rknn.init_runtime('rk3588')if ret != 0:print('Init runtime environment failed!')exit(ret)print('done')# Set inputsimg = 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')outputs = rknn.inference(inputs=[img])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 processinput0_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)rknn.release()