File size: 1,588 Bytes
aec7a65 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 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 |