??雖然車牌識別技術很成熟了，但完全沒有接觸過。一直想搞一下、整一下、試一下、折騰一下，工作之余找了一個簡單的例子入個門。本博客簡單記錄一下 LPRNet 車牌識別部署 rk3588流程，訓練參考 LPRNet 官方代碼。
1、導出onnx
??導出onnx很容易，在推理時加入保存onnx代碼，但用onnx推理時發現推理失敗，是有算子onnx推理時不支持，看了一下不支持的操作 nn.MaxPool3d() ，查了一下資料有等價的方法，用等價方法替換后推理結果是一致的。
import torch.nn as nn
import torchclass 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 xclass small_basic_block(nn.Module):def __init__(self, ch_in, ch_out):super(small_basic_block, self).__init__()self.block = nn.Sequential(nn.Conv2d(ch_in, ch_out // 4, kernel_size=1),nn.ReLU(),nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(3, 1), padding=(1, 0)),nn.ReLU(),nn.Conv2d(ch_out // 4, ch_out // 4, kernel_size=(1, 3), padding=(0, 1)),nn.ReLU(),nn.Conv2d(ch_out // 4, ch_out, kernel_size=1),)def forward(self, x):return self.block(x)class LPRNet(nn.Module):def __init__(self, lpr_max_len, phase, class_num, dropout_rate):super(LPRNet, self).__init__()self.phase = phaseself.lpr_max_len = lpr_max_lenself.class_num = class_numself.backbone = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1), # 0nn.BatchNorm2d(num_features=64),nn.ReLU(), # 2# nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 1, 1)), # 這個可以用MaxPool2d等價nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 1)),small_basic_block(ch_in=64, ch_out=128), # *** 4 ***nn.BatchNorm2d(num_features=128),nn.ReLU(), # 6# nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),maxpool_3d(kernel_size=(1, 3, 3), stride=(2, 1, 2)),small_basic_block(ch_in=64, ch_out=256), # 8nn.BatchNorm2d(num_features=256),nn.ReLU(), # 10small_basic_block(ch_in=256, ch_out=256), # *** 11 ***nn.BatchNorm2d(num_features=256), # 12nn.ReLU(),# nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)), # 14maxpool_3d(kernel_size=(1, 3, 3), stride=(4, 1, 2)), # 14nn.Dropout(dropout_rate),nn.Conv2d(in_channels=64, out_channels=256, kernel_size=(1, 4), stride=1), # 16nn.BatchNorm2d(num_features=256),nn.ReLU(), # 18nn.Dropout(dropout_rate),nn.Conv2d(in_channels=256, out_channels=class_num, kernel_size=(13, 1), stride=1), # 20nn.BatchNorm2d(num_features=class_num),nn.ReLU(), # *** 22 ***)self.container = nn.Sequential(nn.Conv2d(in_channels=448 + self.class_num, out_channels=self.class_num, kernel_size=(1, 1), stride=(1, 1)),# nn.BatchNorm2d(num_features=self.class_num),# nn.ReLU(),# nn.Conv2d(in_channels=self.class_num, out_channels=self.lpr_max_len+1, kernel_size=3, stride=2),# nn.ReLU(),)def forward(self, x):keep_features = list()for i, layer in enumerate(self.backbone.children()):x = layer(x)if i in [2, 6, 13, 22]: # [2, 4, 8, 11, 22]keep_features.append(x)global_context = list()for i, f in enumerate(keep_features):if i in [0, 1]:f = nn.AvgPool2d(kernel_size=5, stride=5)(f)if i in [2]:f = nn.AvgPool2d(kernel_size=(4, 10), stride=(4, 2))(f)f_pow = torch.pow(f, 2)f_mean = torch.mean(f_pow)f = torch.div(f, f_mean)global_context.append(f)x = torch.cat(global_context, 1)x = self.container(x)logits = torch.mean(x, dim=2)return logitsdef build_lprnet(lpr_max_len=8, phase=False, class_num=66, dropout_rate=0.5):Net = LPRNet(lpr_max_len, phase, class_num, dropout_rate)if phase == "train":return Net.train()else:return Net.eval()保存onnx代碼
print("=========== onnx =========== ")
dummy_input = torch.randn(1, 3, 24, 94).cuda()
input_names = ['image']
output_names = ['output']
torch.onnx.export(lprnet, dummy_input, "./weights/LPRNet_model.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=12)
print("======================== convert onnx Finished! .... ")
2 onnx轉換rknn
??onnx轉rknn代碼
# -*- coding: utf-8 -*-import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from math import expimport mathONNX_MODEL = './LPRNet.onnx'
RKNN_MODEL = './LPRNet.rknn'
DATASET = './images_list.txt'QUANTIZE_ON = True'''
CHARS = ['京', '滬', '津', '渝', '冀', '晉', '蒙', '遼', '吉', '黑','蘇', '浙', '皖', '閩', '贛', '魯', '豫', '鄂', '湘', '粵','桂', '瓊', '川', '貴', '云', '藏', '陜', '甘', '青', '寧','新','0', '1', '2', '3', '4', '5', '6', '7', '8', '9','A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K','L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V','W', 'X', 'Y', 'Z', 'I', 'O', '-']'''CHARS = ['BJ', 'SH', 'TJ', 'CQ', 'HB', 'SN', 'NM', 'LN', 'JN', 'HL','JS', 'ZJ', 'AH', 'FJ', 'JX', 'SD', 'HA', 'HB', 'HN', 'GD','GL', 'HI', 'SC', 'GZ', 'YN', 'XZ', 'SX', 'GS', 'QH', 'NX','XJ','0', '1', '2', '3', '4', '5', '6', '7', '8', '9','A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K','L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V','W', 'X', 'Y', 'Z', 'I', 'O', '-']def export_rknn_inference(img):# Create RKNN objectrknn = RKNN(verbose=True)# pre-process configprint('--> Config model')rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588') # mmseprint('done')# Load ONNX modelprint('--> Loading model')ret = rknn.load_onnx(model=ONNX_MODEL, outputs=['output'])if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET, rknn_batch_size=1)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN modelprint('--> Export rknn model')ret = rknn.export_rknn(RKNN_MODEL)if ret != 0:print('Export rknn model failed!')exit(ret)print('done')# Init runtime environmentprint('--> Init runtime environment')ret = rknn.init_runtime()# ret = rknn.init_runtime(target='rk3566')if ret != 0:print('Init runtime environment failed!')exit(ret)print('done')# Inferenceprint('--> Running model')outputs = rknn.inference(inputs=[img])rknn.release()print('done')return outputsif __name__ == '__main__':print('This is main ...')input_w = 94input_h = 24image_path = './test.jpg'origin_image = cv2.imread(image_path)image_height, image_width, images_channels = origin_image.shapeimg = cv2.resize(origin_image, (input_w, input_h), interpolation=cv2.INTER_LINEAR)img = np.expand_dims(img, 0)print(img.shape)preb = export_rknn_inference(img)[0][0]preb_label = []result = []for j in range(preb.shape[1]):preb_label.append(np.argmax(preb[:, j], axis=0))print(preb_label)pre_c = preb_label[0]if pre_c != len(CHARS) - 1:result.append(pre_c)for c in preb_label:if (pre_c == c) or (c == len(CHARS) - 1):if c == len(CHARS) - 1:pre_c = ccontinueresult.append(c)pre_c = cptext = ''for v in result:ptext += CHARS[v]print(ptext)zero_image = np.ones((image_height, image_width, images_channels), dtype=np.uint8) * 255cv2.putText(zero_image, ptext, (0, int(image_height / 2)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2, cv2.LINE_AA)combined_image = np.vstack((origin_image, zero_image))cv2.imwrite('./test_result.jpg', combined_image)
轉換rknn測試結果
說明：由于中文顯示出現亂碼，示例代碼中用拼英簡寫對中文進行了規避
3 部署 rk3588
在rk3588上運行的【完整代碼】
板子上運行結果和時耗。
??模型這么小在rk3588上推理時耗還是比較長的，毫無疑問是模型推理過程中有操作切換到CPU上了。如果對性能要求的比較高，可以針對切換的CPU上的操作進行規避或替換。查看轉換rknn模型log可以知道是那些操作切換到CPU上了。
