fix_yolov5 (wang-xinyu#1181)

* fix_yolov5

* add fix

* add Rex

Author: Rex-LK

yolov5/README.md

@@ -30,6 +30,7 @@ TensorRTx inference code base for [ultralytics/yolov5](https://github.com/ultral
 <a href="https://github.com/triple-Mu"><img src="https://avatars.githubusercontent.com/u/92794867?s=48&v=4" width="40px;" alt=""/></a>
 <a href="https://github.com/xiang-wuu"><img src="https://avatars.githubusercontent.com/u/107029401?s=48&v=4" width="40px;" alt=""/></a>
 <a href="https://github.com/uyolo1314"><img src="https://avatars.githubusercontent.com/u/101853326?s=48&v=4" width="40px;" alt=""/></a>
+<a href="https://github.com/Rex-LK"><img src="https://avatars.githubusercontent.com/u/74702576?s=96&v=4" width="40px;" alt=""/></a>
 
 ## Different versions of yolov5
 

yolov5/yolov5_det_trt.py

@@ -16,7 +16,8 @@
 
 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 @@ def infer(self, raw_image_generator):
         # 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 @@ def post_process(self, output, origin_h, origin_w):
         # 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([])