import cv2 as cv
import numpy as np
import time

import sys

from regex import M
from util import ResizeImage,CropImage,ToTensor,NormalizeImage
from openvino.runtime import Core

model_xml = "IR/handwriting.xml"
model_bin = "IR/handwriting.bin"

onnx_model = "onnx/handwriting.onnx"
image_file = "test/test4.png"

top_k = 10 

# 图像预处理，转为Tensor张量
def preprocess(img):
    resize_op = ResizeImage(resize_short=256)
    img = resize_op(img)
    crop_op = CropImage(size=(224, 224))
    img = crop_op(img)
    img_mean = [0.485, 0.456, 0.406]
    img_std = [0.229, 0.224, 0.225]
    img_scale = 1.0 / 255.0
    normalize_op = NormalizeImage(
        scale=img_scale, mean=img_mean, std=img_std)
    img = normalize_op(img)
    tensor_op = ToTensor()
    img = tensor_op(img)
    return img

#标签文件
label_file = "labels.txt"

# 读取标签文件
def readLabels():
    index = 0
    labels=[]
    with open(label_file) as file_obj:
        for line in file_obj:
            labels.append(line.strip())
            index=index+1
    return labels

labels = readLabels()

# load model
ie = Core()
model = ie.read_model(model=onnx_model)
compiled_model = ie.compile_model(model=model, device_name="CPU")
output_layer = next(iter(compiled_model.outputs))
input_layer = next(iter(compiled_model.inputs))

# prepare image
frame = cv.imread(image_file)[:,:,::-1]
# frame = cv.resize(frame, (64,64))
# deal
# n, c, h, w = net.inputs[input_blob].shape
# initial_h, initial_w, channels = frame.shape
#按照AI模型要求放缩图片
image = cv.resize(frame, (224, 224))  
image = preprocess(image)
input_image = np.expand_dims(image.transpose(0,1,2),0)

# do inference
result_infer = compiled_model(inputs=[input_image])[output_layer]
print(result_infer[0])
result_index = result_infer[0].argsort()[-10:][::-1]
print(result_index)
for index in result_index:
    print(labels[index])