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)