predict-ov2.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. import cv2 as cv
  2. import numpy as np
  3. import time
  4. import sys
  5. from regex import M
  6. from util import ResizeImage,CropImage,ToTensor,NormalizeImage
  7. from openvino.runtime import Core
  8. model_xml = "IR/handwriting.xml"
  9. model_bin = "IR/handwriting.bin"
  10. onnx_model = "onnx/handwriting.onnx"
  11. image_file = "test/test4.png"
  12. top_k = 10
  13. # 图像预处理,转为Tensor张量
  14. def preprocess(img):
  15. resize_op = ResizeImage(resize_short=256)
  16. img = resize_op(img)
  17. crop_op = CropImage(size=(224, 224))
  18. img = crop_op(img)
  19. img_mean = [0.485, 0.456, 0.406]
  20. img_std = [0.229, 0.224, 0.225]
  21. img_scale = 1.0 / 255.0
  22. normalize_op = NormalizeImage(
  23. scale=img_scale, mean=img_mean, std=img_std)
  24. img = normalize_op(img)
  25. tensor_op = ToTensor()
  26. img = tensor_op(img)
  27. return img
  28. #标签文件
  29. label_file = "labels.txt"
  30. # 读取标签文件
  31. def readLabels():
  32. index = 0
  33. labels=[]
  34. with open(label_file) as file_obj:
  35. for line in file_obj:
  36. labels.append(line.strip())
  37. index=index+1
  38. return labels
  39. labels = readLabels()
  40. # load model
  41. ie = Core()
  42. model = ie.read_model(model=onnx_model)
  43. compiled_model = ie.compile_model(model=model, device_name="CPU")
  44. output_layer = next(iter(compiled_model.outputs))
  45. input_layer = next(iter(compiled_model.inputs))
  46. # prepare image
  47. frame = cv.imread(image_file)[:,:,::-1]
  48. # frame = cv.resize(frame, (64,64))
  49. # deal
  50. # n, c, h, w = net.inputs[input_blob].shape
  51. # initial_h, initial_w, channels = frame.shape
  52. #按照AI模型要求放缩图片
  53. image = cv.resize(frame, (224, 224))
  54. image = preprocess(image)
  55. input_image = np.expand_dims(image.transpose(0,1,2),0)
  56. # do inference
  57. result_infer = compiled_model(inputs=[input_image])[output_layer]
  58. print(result_infer[0])
  59. result_index = result_infer[0].argsort()[-10:][::-1]
  60. print(result_index)
  61. for index in result_index:
  62. print(labels[index])