import time
|
import cv2
|
import numpy as np
|
import yaml
|
from openvino.runtime import Core
|
|
|
class IDENTIFIER:
|
|
def __init__(self, path):
|
# Initialize model
|
self.initialize_model(path)
|
|
def __call__(self, image):
|
return self.idengify(image)
|
def read_config(self, path):
|
file_path = path+'/metadata.yaml'
|
with open(file_path, 'r', encoding="utf-8") as file:
|
config = yaml.safe_load(file)
|
return config
|
def initialize_model(self, path):
|
model_path = path + '/best.xml'
|
# Initialize OpenVINO Runtime
|
self.core = Core()
|
# Load the model
|
self.model = self.core.read_model(model=model_path)
|
# Compile the model
|
self.compiled_model = self.core.compile_model(model=self.model, device_name="CPU")
|
# Get input and output layers
|
self.input_layer = self.compiled_model.input(0)
|
N,C,self.input_width,self.input_height = self.input_layer.shape
|
|
self.output_layer = self.compiled_model.output(0)
|
# Get class names
|
self.class_names = CLASSES = self.read_config(path)['names']
|
# Get model info
|
# self.get_input_details()
|
# self.get_output_details()
|
|
def idengify(self, image):
|
input_tensor = self.prepare_input(image)
|
|
# Perform inference on the image
|
outputs = self.inference(input_tensor)
|
|
|
|
return outputs
|
|
def prepare_input(self, image):
|
self.img_height, self.img_width = image.shape[:2]
|
|
input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
# Resize图片不要直接使用resize，需要按比例缩放，空白区域填空纯色即可
|
# input_img = self.ratioresize(input_img)
|
# 启用中心裁剪
|
input_img = self.center_crop(input_img)
|
|
# Scale input pixel values to 0 to 1
|
input_img = input_img / 255.0
|
input_img = input_img.transpose(2, 0, 1)
|
input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)
|
|
return input_tensor
|
|
def inference(self, input_tensor):
|
ir = self.compiled_model.create_infer_request()
|
outs = ir.infer(input_tensor)[self.output_layer]
|
|
# print(f"Inference time: {(time.perf_counter() - start)*1000:.2f} ms")
|
return outs
|
|
|
|
# def get_input_details(self):
|
# model_inputs = self.session.get_inputs()
|
# self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
|
#
|
# self.input_shape = model_inputs[0].shape
|
# self.input_height = self.input_shape[2]
|
# self.input_width = self.input_shape[3]
|
|
# def get_output_details(self):
|
# model_outputs = self.session.get_outputs()
|
# self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
|
|
# 等比例缩放图片
|
def ratioresize(self, im, color=114):
|
shape = im.shape[:2]
|
new_h, new_w = self.input_height, self.input_width
|
padded_img = np.ones((new_h, new_w, 3), dtype=np.uint8) * color
|
|
# Scale ratio (new / old)
|
r = min(new_h / shape[0], new_w / shape[1])
|
|
# Compute padding
|
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
|
|
if shape[::-1] != new_unpad:
|
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
|
|
padded_img[: new_unpad[1], : new_unpad[0]] = im
|
padded_img = np.ascontiguousarray(padded_img)
|
return padded_img
|
|
def center_crop(self, img):
|
# 新增中心裁剪方法
|
h, w = img.shape[:2]
|
desired_h = self.input_height
|
desired_w = self.input_width
|
|
# 如果图片尺寸大于目标尺寸，则进行中心裁剪
|
if h > desired_h and w > desired_w:
|
start_y = (h - desired_h) // 2
|
start_x = (w - desired_w) // 2
|
end_y = start_y + desired_h
|
end_x = start_x + desired_w
|
return img[start_y:end_y, start_x:end_x]
|
else:
|
# 否则进行缩放
|
return self.ratioresize(img)
|
