haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update yolov8_det_trt.py

Browse Source
This commit is contained in:
Wang Xinyu 2023-11-03 16:18:36 +08:00 committed by GitHub
parent 06d3a40a78
commit f39f2c1433
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
yolov8/yolov8_det_trt.py
View File

@ -167,7 +167,7 @@ class YoLov8TRT(object):
# Do postprocess
for i in range(self.batch_size):
result_boxes, result_scores, result_classid = self.post_process(
output[i * 6001: (i + 1) * 6001], batch_origin_h[i], batch_origin_w[i]
output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i]
)
# Draw rectangles and labels on the original image
for j in range(len(result_boxes)):
@ -290,7 +290,7 @@ class YoLov8TRT(object):
# Get the num of boxes detected
num = int(output[0])
# Reshape to a two dimentional ndarray
pred = np.reshape(output[1:], (-1, 6))[:num, :]
pred = np.reshape(output[1:], (-1, 38))[:num, :]
# Do nms
boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
result_boxes = boxes[:, :4] if len(boxes) else np.array([])