initial commit

yolo_server/infer.py

@@ -0,0 +1,65 @@
+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')
+