From 883d0cc566db09b9f9e356cbb53eb5c5a5efc10d Mon Sep 17 00:00:00 2001
From: acai66 <1779864536@qq.com>
Date: Tue, 28 Dec 2021 17:58:39 +0800
Subject: [PATCH] Fix ValueError: cannot select an axis to squeeze out which
 has size not equal to one.

---
 utils/general.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils/general.py b/utils/general.py
index 8d46824..0a5252d 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -742,7 +742,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
         # i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
         
         i = cv2.dnn.NMSBoxesRotated(boxes_for_cv2_nms, scores_for_cv2_nms, conf_thres, iou_thres)
-        i = np.squeeze(i, axis=-1)
+        if i.ndim > 1:
+            i = np.squeeze(i, axis=-1)
         
         if i.shape[0] > max_det:  # limit detections
             i = i[:max_det]