From d6a0b784c431bc62a3167ba48d112b7670845ccf Mon Sep 17 00:00:00 2001
From: Rex <74702576+Rex-LK@users.noreply.github.com>
Date: Wed, 21 Dec 2022 19:40:23 +0800
Subject: [PATCH] fix_yolov5 (#1181)
* fix_yolov5
* add fix
* add Rex
---
yolov5/README.md | 1 +
yolov5/yolov5_det_trt.py | 8 +++++---
2 files changed, 6 insertions(+), 3 deletions(-)
diff --git a/yolov5/README.md b/yolov5/README.md
index 94113e4..c4db012 100644
--- a/yolov5/README.md
+++ b/yolov5/README.md
@@ -30,6 +30,7 @@ TensorRTx inference code base for [ultralytics/yolov5](https://github.com/ultral
+
## Different versions of yolov5
diff --git a/yolov5/yolov5_det_trt.py b/yolov5/yolov5_det_trt.py
index b2b03d9..bed0ba4 100644
--- a/yolov5/yolov5_det_trt.py
+++ b/yolov5/yolov5_det_trt.py
@@ -16,7 +16,8 @@ import tensorrt as trt
CONF_THRESH = 0.5
IOU_THRESHOLD = 0.4
-
+LEN_ALL_RESULT = 38001
+LEN_ONE_RESULT = 38
def get_img_path_batches(batch_size, img_dir):
ret = []
@@ -166,7 +167,7 @@ class YoLov5TRT(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 * LEN_ALL_RESULT: (i + 1) * LEN_ALL_RESULT], batch_origin_h[i], batch_origin_w[i]
)
# Draw rectangles and labels on the original image
for j in range(len(result_boxes)):
@@ -289,7 +290,8 @@ class YoLov5TRT(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, LEN_ONE_RESULT))[:num, :]
+ pred = pred[:, :6]
# 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([])