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)