import paddle.inference as paddle_infer
import numpy as np
import cv2
class Yolo_model_infer:
    def __init__(self, model_dir="./ppyoloe_plus_crn_t_auxhead_320_300e_coco", target_size=[320, 320]):
        # 初始化 paddle 推理
        self.model_dir = model_dir
        self.config = paddle_infer.Config(model_dir + "/model.pdmodel", model_dir + "/model.pdiparams")
        self.config.disable_glog_info()
        self.config.enable_use_gpu(500, 0)
        self.config.enable_memory_optim()
        self.config.switch_ir_optim()
        self.config.switch_use_feed_fetch_ops(False)
        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
        origin_shape = (240, 320)
        resize_h, resize_w = self.target_size
        self.im_scale_y = resize_h / float(origin_shape[0])
        self.im_scale_x = resize_w / float(origin_shape[1])
        self.scale_info = np.array([[self.im_scale_y, self.im_scale_x]]).astype('float32')
    def infer(self,src) -> np.ndarray:
        image = self.preprocess(src)
        self.input_handle.copy_from_cpu(image)
        self.input_handle1.copy_from_cpu(self.scale_info)
        self.predictor.run()
        results = self.output_handle.copy_to_cpu()
        return results
    def preprocess(self,src):
        # resize
        # keep_ratio=0
        img = cv2.resize(
            src,
            None,
            None,
            fx=self.im_scale_x,
            fy=self.im_scale_y,
            interpolation=2)
        # NormalizeImage
        img = img.astype(np.float32, copy=False)
        scale = 1.0 / 255.0
        img *= scale
        # Permute
        img = img.transpose((2, 0, 1))
        img = np.array((img, ))
        # .astype('float32')
        return img