Create utils.py

utils.py

@@ -0,0 +1,44 @@
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+def preprocess(image, size=(640, 640)):
+    h, w = image.shape[:2]
+    max_size = max(h, w)
+    scale_factor = size[0] / max_size
+    pad_h = (max_size - h) // 2
+    pad_w = (max_size - w) // 2
+    pad_image = np.zeros((max_size, max_size, 3), dtype=image.dtype)
+    pad_image[pad_h:h + pad_h, pad_w:w + pad_w] = image
+    image = cv2.resize(pad_image, size,
+                       interpolation=cv2.INTER_LINEAR).astype('float32')
+    image /= 255.0
+    image = image[None]
+    return image, scale_factor, (pad_h, pad_w)
+
+def inference(ort_session,
+              ori_image,
+              size=(640, 640),
+              **kwargs):
+    # normal export
+    # with NMS and postprocessing
+    h, w = ori_image.shape[:2]
+    image, scale_factor, pad_param = preprocess(ori_image[:, :, [2, 1, 0]],
+                                                size)
+    input_ort = ort.OrtValue.ortvalue_from_numpy(image.transpose((0, 3, 1, 2)))
+    results = ort_session.run(["num_dets", "labels", "scores", "boxes"],
+                              {"images": input_ort})
+    num_dets, labels, scores, bboxes = results
+    num_dets = num_dets[0][0]
+    labels = labels[0, :num_dets]
+    scores = scores[0, :num_dets]
+    bboxes = bboxes[0, :num_dets]
+
+    bboxes -= np.array(
+        [pad_param[1], pad_param[0], pad_param[1], pad_param[0]])
+    bboxes /= scale_factor
+    bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0, w)
+    bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, h)
+    bboxes = bboxes.round().astype('int')
+
+    return labels, scores, bboxes