LPRNet_migraphx_infer.py

# -*- coding: utf-8 -*-
"""
MIGraphX示例程序
"""
import cv2
import numpy as np
import migraphx

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', '-'
         ]

def LPRNetPreprocess(image):
    img = cv2.imread(image)
    img = cv2.resize(img, (94, 24)).astype('float32')
    img -= 127.5
    img *= 0.0078125
    img = np.expand_dims(img.transpose(2, 0, 1), 0)
    return img

def LPRNetPostprocess(infer_res):
    preb_label = []
    for j in range(infer_res.shape[1]):
        preb_label.append(np.argmax(infer_res[:, j], axis=0))
    no_repeat_blank_label = []
    pre_c = preb_label[0]
    if pre_c != len(CHARS) - 1:
        no_repeat_blank_label.append(pre_c)
    for c in preb_label:  # dropout repeate label and blank label
        if (pre_c == c) or (c == len(CHARS) - 1):
            if c == len(CHARS) - 1:
                pre_c = c
            continue
        no_repeat_blank_label.append(c)
        pre_c = c
    result = ''.join(list(map(lambda x: CHARS[x], no_repeat_blank_label)))
    return result

def LPRNetInference(model_name, imgs):
    img = LPRNetPreprocess(imgs)
    
    # 加载模型
    if model_name[-3:] == 'mxr':
        model = migraphx.load(model_name)
    else:
        print('convert onnx to mxr...')
        model = migraphx.parse_onnx(model_name)
        model.compile(t=migraphx.get_target("gpu"),device_id=0) # device_id: 设置GPU设备，默认为0号设备(>=1.2版本中支持)
        migraphx.save(model, 'model/LPRNet.mxr')

    inputName=model.get_parameter_names()[0]
    inputShape=model.get_parameter_shapes()[inputName].lens()
    print("inputName:{0} \ninputShape:{1}".format(inputName,inputShape))

    results = model.run({inputName: migraphx.argument(img)})

    result = LPRNetPostprocess(np.array(results[0]))
    return result

if __name__ == '__main__':
    # model_name = 'model/LPRNet.onnx'
    model_name = 'model/LPRNet.mxr'
    image = 'imgs/川JK0707.jpg'
    InferRes = LPRNetInference(model_name, image)
    print(image, 'Inference Result:', InferRes)