predict.py 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. from openvino.inference_engine import IECore
  2. import cv2 as cv
  3. import numpy as np
  4. import time
  5. import sys
  6. from util import ResizeImage,CropImage,ToTensor,NormalizeImage
  7. model_xml = "IR/handwriting.xml"
  8. model_bin = "IR/handwriting.bin"
  9. onnx_model = "onnx/handwriting.onnx"
  10. image_file = "test/test3.jpg"
  11. top_k = 10
  12. # 图像预处理,转为Tensor张量
  13. def preprocess(img):
  14. resize_op = ResizeImage(resize_short=256)
  15. img = resize_op(img)
  16. crop_op = CropImage(size=(224, 224))
  17. img = crop_op(img)
  18. img_mean = [0.485, 0.456, 0.406]
  19. img_std = [0.229, 0.224, 0.225]
  20. img_scale = 1.0 / 255.0
  21. normalize_op = NormalizeImage(
  22. scale=img_scale, mean=img_mean, std=img_std)
  23. img = normalize_op(img)
  24. tensor_op = ToTensor()
  25. img = tensor_op(img)
  26. return img
  27. #标签文件
  28. label_file = "labels.txt"
  29. # 读取标签文件
  30. def readLabels():
  31. index = 0
  32. labels=[]
  33. with open(label_file) as file_obj:
  34. for line in file_obj:
  35. labels.append(line.strip())
  36. index=index+1
  37. return labels
  38. labels = readLabels()
  39. DEVICE = 'CPU'
  40. #初始化插件,输出插件版本号
  41. ie = IECore()
  42. devices = ie.available_devices
  43. print(devices)
  44. ver = ie.get_versions(DEVICE)[DEVICE]
  45. print("{descr}: {maj}.{min}.{num}".format(descr=ver.description, maj=ver.major, min=ver.minor, num=ver.build_number))
  46. #读取IR模型文件
  47. # net = ie.read_network(model=model_xml, weights=model_bin)
  48. net = ie.read_network(model=onnx_model)
  49. #准备输入输出张量MYRIAD
  50. print("Preparing input blobs")
  51. input_blob = next(iter(net.inputs))
  52. out_blob = next(iter(net.outputs))
  53. net.batch_size = 1
  54. #载入模型到AI推断计算设备
  55. print("Loading IR to the plugin...")
  56. exec_net = ie.load_network(network=net, num_requests=1, device_name=DEVICE)
  57. frame = cv.imread(image_file)[:,:,::-1]
  58. # frame = cv.resize(frame, (64,64))
  59. # deal
  60. n, c, h, w = net.inputs[input_blob].shape
  61. initial_h, initial_w, channels = frame.shape
  62. #按照AI模型要求放缩图片
  63. image = cv.resize(frame, (w, h))
  64. image = preprocess(image)
  65. print("Batch size is {}".format(n))
  66. #执行推断计算
  67. print("Starting inference in synchronous mode")
  68. start = time.time()
  69. res = exec_net.infer(inputs={input_blob: image})
  70. end = time.time()
  71. print("Infer Time:{}ms".format((end-start)*1000))
  72. # 处理输出
  73. print("Processing output blob")
  74. res = res[out_blob]
  75. result = zip(res[0].tolist(), labels)
  76. result2 = sorted(result, reverse=True)
  77. candidate = list(result2[0:top_k])
  78. for c in candidate:
  79. print(c[1])