import time
|
import cv2
|
import numpy as np
|
import onnxruntime
|
|
|
class IDENTIFIER:
|
|
def __init__(self, path):
|
# Initialize model
|
self.initialize_model(path)
|
|
def __call__(self, image):
|
return self.idengify(image)
|
|
def initialize_model(self, path):
|
self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
self.class_names = eval(self.session.get_modelmeta().custom_metadata_map['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)
|
|
self.herb_probabilities = outputs[0]
|
|
return self.herb_probabilities
|
|
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):
|
start = time.perf_counter()
|
outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor})
|
|
# print(f"Inference time: {(time.perf_counter() - start)*1000:.2f} ms")
|
return outputs
|
|
|
|
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)
|
