LPRNet_ORT_infer.py

import onnxruntime as ort
import cv2
import numpy as np
import argparse
import os
import time

print('Runing Based On:', ort.get_device())

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 = np.argmax(infer_res, 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(args):
    if ort.get_device() == "GPU-MIGRAPHX":
        sess = ort.InferenceSession(args.model, providers=['ROCMExecutionProvider'],) #DCU版本
    else:
        sess = ort.InferenceSession(args.model, providers=['CPUExecutionProvider']) # CPU版本

    if os.path.isdir(args.imgpath):
        images = os.listdir(args.imgpath)
        Tp = 0
        Tn_1 = 0
        Tn_2 = 0
        time1 = time.perf_counter()
        for image in images:
            img = LPRNetPreprocess(os.path.join(args.imgpath, image))
            intput = sess.get_inputs()[0].shape
            preb = sess.run(None, input_feed={sess.get_inputs()[0].name: img})[0]
            result = LPRNetPostprocess(preb)
            if result == image[:-4]:
                Tp += 1
            elif len(result) != len(image[:-4]):
                Tn_1 += 1
            else:
                Tn_2 += 1
            print(image + ' Inference Result:', result)
        time2 = time.perf_counter() 
        Acc = Tp * 1.0 / (Tp + Tn_1 + Tn_2)
        print("[Info] Test Accuracy: {} [{}:{}:{}:{}]".format(Acc, Tp, Tn_1, Tn_2, (Tp+Tn_1+Tn_2)))
        print("[Info] Test Speed: {}s 1/{}]".format((time2 - time1) / len(images), len(images)))     
    else:
        img = LPRNetPreprocess(args.imgpath)
        intput = sess.get_inputs()[0].shape
        preb = sess.run(None, input_feed={sess.get_inputs()[0].name: img})[0]
        result = LPRNetPostprocess(preb)
        print('Inference Result:', result)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='parameters to vaildate net')
    parser.add_argument('--model', default='model/LPRNet.onnx', help='model path to vaildate')
    parser.add_argument('--imgpath', default='imgs', help='the image path')
    args = parser.parse_args()

    LPRNetInference(args)