|
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): |
|
|
|
|
|
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 |
