if crnn_vertical_model_path is not None:
crnn_vertical_net = LiteCrnn(32,
1,
len(alphabet) + 1,
nh,
n_rnn=2,
leakyRelu=False,
lstmFlag=True)
crnn_vertical_handle = CRNNHandle(crnn_vertical_model_path,
crnn_vertical_net,
gpu_id=GPU_ID)
assert angle_type in ["shufflenetv2_05"]
if angle_type == "shufflenetv2_05":
angle_net = shufflenet_v2_x0_5(num_classes=len(lable_map_dict),
pretrained=False)
angle_handle = AangleClassHandle(angle_model_path, angle_net, gpu_id=GPU_ID)
def crnnRec(im,
rects_re,
leftAdjust=False,
rightAdjust=False,
alph=0.2,
f=1.0):
"""
crnn模型，ocr识别
@@model,
@@converter,
@@im:Array
import sys
import mxnet as mx
sys.path.insert(0, "..")
from crnn import LiteCrnn
from utils0nnx import load_model
import torch
from angle_class import shufflenet_v2_x0_5
from crnn.keys import alphabetChinese as alphabet
nh = 256
net = shufflenet_v2_x0_5(num_classes=4, pretrained=False)
net = load_model(net, '../models/shufflenetv2_05.pth')
net.eval()
print('Finished loading model!')
print(net)
device = torch.device("cpu")
net = net.to(device)
##################export###############
output_onnx = 'shufflenetv2_05.onnx'
print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
input_names = ["input"]
# output_names = ["hm" , "wh" , "reg"]
output_names = ["out"]
inputs = torch.randn(1, 3, 48, 196).to(device)
torch_out = torch.onnx._export(net,
inputs,
output_onnx,
export_params=True,