Exemplo n.º 1
0
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,