-
Couldn't load subscription status.
- Fork 1.6k
Closed
Description
I found how to inference a trt yolov9 model using this detect function:
def detect(self, bgr_img):
## Padded resize
h, w, _ = bgr_img.shape
scale = min(self.imgsz[0]/w, self.imgsz[1]/h)
inp = np.zeros((self.imgsz[1], self.imgsz[0], 3), dtype = np.float32)
nh = int(scale * h)
nw = int(scale * w)
inp[: nh, :nw, :] = cv2.resize(cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB), (nw, nh))
inp = inp.astype('float32') / 255.0 # 0 - 255 to 0.0 - 1.0
inp = np.expand_dims(inp.transpose(2, 0, 1), 0)
## Inference
t1 = time.time()
num_detection, nmsed_bboxes, nmsed_scores, nmsed_classes = self.model.run(inp)
t2 = time.time()
## Apply NMS
num_detection = num_detection[0][0]
nmsed_bboxes = nmsed_bboxes[0]
nmsed_scores = nmsed_scores[0]
nmsed_classes = nmsed_classes[0]
print('Detected {} object(s)'.format(num_detection))
# Rescale boxes from img_size to im0 size
_, _, height, width = inp.shape
h, w, _ = bgr_img.shape
nmsed_bboxes[:, 0] /= scale
nmsed_bboxes[:, 1] /= scale
nmsed_bboxes[:, 2] /= scale
nmsed_bboxes[:, 3] /= scale
visualize_img = bgr_img.copy()
for ix in range(num_detection): # x1, y1, x2, y2 in pixel format
cls = int(nmsed_classes[ix])
label = '%s %.2f' % (self.names[cls], nmsed_scores[ix])
x1, y1, x2, y2 = nmsed_bboxes[ix]
cv2.rectangle(visualize_img, (int(x1), int(y1)), (int(x2), int(y2)), self.colors[int(cls)], 2)
cv2.putText(visualize_img, label, (int(x1), int(y1-10)), cv2.FONT_HERSHEY_SIMPLEX, 1, self.colors[int(cls)], 2, cv2.LINE_AA)
cv2.imwrite('result.jpg', visualize_img)
return `visualize_img`
But I can't find out how to do the same for an ONNX model. Has someone already made something like this, or should I do it and post it? Thank you.
Metadata
Metadata
Assignees
Labels
No labels