From 7b79de466c7ac2fcf179e65c2fa4718107f236f9 Mon Sep 17 00:00:00 2001
From: Rex <74702576+Rex-LK@users.noreply.github.com>
Date: Tue, 3 Jan 2023 15:58:16 +0800
Subject: [PATCH] fix a problem in warmUp (#1194)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Int process of warmUp，the pred is None, resulting in the error of "masks = self.sigmoid((result_proto_coef @ result_proto_masks.astype(np.float32).reshape(c, -1))).reshape(-1, mh, mw)"
---
 yolov5/yolov5_seg_trt.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/yolov5/yolov5_seg_trt.py b/yolov5/yolov5_seg_trt.py
index bdc0465..b0ac307 100644
--- a/yolov5/yolov5_seg_trt.py
+++ b/yolov5/yolov5_seg_trt.py
@@ -181,6 +181,8 @@ class YoLov5TRT(object):
             result_boxes, result_scores, result_classid, result_proto_coef = self.post_process(
                 output_bbox[i * self.det_output_length: (i + 1) * self.det_output_length], batch_origin_h[i], batch_origin_w[i]
             )
+            if result_proto_coef.shape[0] == 0:
+                continue
             result_masks = self.process_mask(output_proto_mask, result_proto_coef, result_boxes, batch_origin_h[i], batch_origin_w[i])
 
             # Draw masks on  the original image