车牌识别部署瑞芯微RK3588

因为国内关于rockchip方面的文档不是很多,所以在此记录一下车牌识别模型转化以及部署的过程:

数据集:

cpdd,github上寻找到数据集

链接: https://pan.baidu.com/s/1SeqcoBBDqOYVBvpOTZy7tQ 提取码: 46wk

模型:

yolov5,crnn(lprnet_rknn模型量化失败)

车牌检测

yolov5模型不过多赘述,网络上面开源的很多,转化过程我分两步,pytorch转onnx,再onnx转rknn.
yolov5开源代码中有export.py文件用于输出onnx模型

python3 export.py —weights weights/yolov5s.pt —opset 12

在onnx转rknn模型时可能会出现报错,需要修改opset参数(目前没有遇到这种情况).
rknn_toolkit2官方提供了python3.6和python3.8的package,如果pc没有这两个版本,就需要按转对应版本.
官方有提供onnx转rknn模型的代码(rknn-toolkit2-1.4.0/examples/onnx/yolov5/):

# Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[1, 1, 1]], target_platform='rk3588')
    print('done')

    # Load ONNX model
    print('--> Loading model')
    # ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['326', '384', '442'])
    ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['343', '360', '377'])
    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
    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()
    # ret = rknn.init_runtime(target='rk3588',device_id='85b9790400346ffd', async_mode=True)
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')
    img = cv2.imread(IMG_PATH)
    img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

rknn的接口定义在rknn-toolkit2-1.4.0/doc/Rockchip_User_Guide_RKNN_Toolkit2_CN-1.4.0.pdf文件中.其中主要需要修改的部分一般有rknn.config,rknn.load_onnx,init_runtime,rknn.inference.

load_onnx函数中outputs需要设置;可以通过Netron来查看onnx模型的输出节点,主要找到这三个输出节点:

init_runtime函数用来模型结果查看:

其中device_id通过adb devices命令查看添加,target为目标平台。
inference函数得到模型输出。

车牌识别

之前使用的lprnet来做车牌识别的工作,,用了50000张图片,训练30个epochs能够达到92%的准确率.但在转换模型时出现了问题,由于torch.onnx.export()对于maxpool3d不支持,需要将maxpool3d转化为3个maxpool2d来适应onnx转换。

class maxpool_3d(nn.Module):
    def __init__(self, kernel_size, stride):
        super(maxpool_3d, self).__init__()
        assert(len(kernel_size)==3 and len(stride)==3)
        kernel_size2d1 = kernel_size[-2:]
        stride2d1 = stride[-2:]
        kernel_size2d2 = (kernel_size[0],kernel_size[0])
        stride2d2 = (kernel_size[0], stride[0])
        self.maxpool1 = nn.MaxPool2d(kernel_size=kernel_size2d1, stride=stride2d1)
        self.maxpool2 = nn.MaxPool2d(kernel_size=kernel_size2d2, stride=stride2d2)

    def forward(self,x):
        x = self.maxpool1(x)
        x = x.transpose(1,3)
        x = self.maxpool2(x)
        x = x.transpose(1,3)
        return x

使用后在torch.onnx.export()步骤不会出现报错了,但是我在使用onnxruntime进行结果复现时出现了问题,结果不对等,准确率掉到了80%,同时在转化为rknn模型后准确率掉到了75%.

在进行了一系列数据读取的不同操作之后,选择了放弃,结果始终无法提升(具体原因由于个人能力问题,无法解释).

后续选择了更换模型,换成了crnn网络,使用同样方式进行了模型转换,此次结果onnx模型和pytorch模型结果对比一致(想想还是没道理,觉得应该是网络结构有问题,但是pytorch测试时识别是没有问题的).

模型转换与yolov5是一样的方式:

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> Config model')
    # rknn.config(mean_values=[[127.5, 127.5, 127.5]], std_values=[[255, 255, 255]], target_platform='rk3588')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[1, 1, 1]], target_platform='rk3588')
    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=False)
    if ret != 0:
        print('Build model 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()
    # ret = rknn.init_runtime(target='rk3588', device_id='85b9790400346ffd', async_mode=True)
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

不过在rknn.inference函数上需要进行修改,如果想要与pytorch一样的前处理方式,在进行transpose()转换维度时,

    img = cv2.resize(img, (168, 48))
    img = img.astype(np.float32)
    img = (img / 255 - mean_value) / std_value  # 归一化 减均值 除标准差
    img = img.transpose(2, 0, 1)  # h,w,c 转为 c,h,w
    img = img.reshape(1, *img.shape)  # channel,height,width转为batch,channel,height,channel

需要将data_format参数修改为nchw模式,但其实没有必要进行transpose操作(后续会说)

outputs = rknn.inference(inputs=[img], data_format='nchw')[0]

后续进行结果的解析就与yolov5不太一样了,因为crnn是单输出形式,与yolov5的三输出形式不一样,所以只需要进行argmax操作就可以了,然后就配合classes类别输出结果就可以了.

到此部分都是在pc端进行,都是对结果的模拟,并不能直接放置到rk板子上运行。

模型部署

官方github下有rknpu2项目代码,此部分是关于模型部署.

不过我是基于rknn_toolkit2下rknn_toolkit_lite2完成的部署(因为本人C++技术堪称没有,只能依靠python方式啦),rknn_toolkit_lite2需要python3.7或python3.9部署.(进行这些操作之前需要在板端进行rknn_server和librknnrt.so的更新,具体可以看rknpu2/rknn_server_proxy.md,依照自己方式更新,我的是linux.)

rknn_toolkit_lite2与pc端rknn_toolkit2没什么太大区别,模型的家在初始化基本上是一样的,就是将rknn更换成rknn_lite,函数也都大差不差,更细节的可以看api接口文档.

rknn_lite = RKNNLite()

    print('-->loading model : ' + RKNN_MODEL)
    ret = rknn_lite.load_rknn(RKNN_MODEL)
    # Init runtime environment
    print('--> Init runtime environment')
    # ret = rknn.init_runtime()
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO)
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

结果预测与rknn也差不多,区别是在rknn_toolkit_lite下inference中data_format参数只有’nhwc’一种数据排列方式,所以在数据前处理部分只需要将transpose这一步操作取消掉就可以了。

    outputs = rknn.inference(inputs=[img])[0]

后续处理就没什么区别了,只需要结合自己的业务进行代码修改就可以了。