fix a problem in warmUp (#1194)

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)"
This commit is contained in:
Rex 2023-01-03 15:58:16 +08:00 committed by GitHub
parent 3c0f62e82f
commit 7b79de466c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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