import paddle.inference as paddle_infer
import numpy as np
import cv2
class Yolo_model_infer:
    def __init__(self, model_dir="./yolo_model", target_size=[640, 640]):
        # 初始化 paddle 推理
        self.model_dir = model_dir
        self.config = paddle_infer.Config(model_dir + "/model.pdmodel", model_dir + "/model.pdiparams")
        self.config.enable_memory_optim()
        self.config.switch_ir_optim()
        self.config.enable_use_gpu(1000, 0)
        
        self.predictor = paddle_infer.create_predictor(self.config)
        self.input_names = self.predictor.get_input_names()
        self.input_handle = self.predictor.get_input_handle(self.input_names[0])
        self.input_handle1 = self.predictor.get_input_handle(self.input_names[1])
        self.output_names = self.predictor.get_output_names()
        self.output_handle = self.predictor.get_output_handle(self.output_names[0])

        self.target_size = target_size
        self.fill_value = [114.0, 114.0, 114.0]
    def infer(self,src) -> np.ndarray:
        image, scale_info = self.preprocess(src)
        
        self.input_handle.copy_from_cpu(image)
        self.input_handle1.copy_from_cpu(scale_info)
        self.predictor.run()
        results = self.output_handle.copy_to_cpu()
        return results
    def preprocess(self,src):
        # resize
        origin_shape = src.shape[:2]
        # keep_ratio==1
        im_size_min = np.min(origin_shape)
        im_size_max = np.max(origin_shape)
        target_size_min = np.min(self.target_size)
        target_size_max = np.max(self.target_size)
        im_scale = float(target_size_min) / float(im_size_min)
        if np.round(im_scale * im_size_max) > target_size_max:
            im_scale = float(target_size_max) / float(im_size_max)
        im_scale_x = im_scale
        im_scale_y = im_scale

        img = cv2.resize(
            src,
            None,
            None,
            fx=im_scale_x,
            fy=im_scale_y,
            interpolation=1)
        # pad
        # pad = Pad((640, 640))
        # img = pad(img)
        im_h, im_w = img.shape[:2]
        h, w = self.target_size
        canvas = np.ones((h, w, 3), dtype=np.float32)
        canvas *= np.array(self.fill_value, dtype=np.float32)
        canvas[0:im_h, 0:im_w, :] = img.astype(np.float32)
        img = canvas

        # Permute
        img = img.transpose((2, 0, 1)).copy()
        img = np.array((img, )).astype('float32')
        return img, np.array([im_scale_y, im_scale_x]).astype('float32')