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 = []
    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(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)
        count = 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]:
                count += 1
            print('Inference Result:', result)
        time2 = time.perf_counter()
        print('accuracy rate:', count / len(images))
        print('average time', (time2 - time1)/count*1000)        
    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)