Fix ValueError: cannot select an axis to squeeze out which has size n…

…ot equal to one.
acai66 · Dec 28, 2021 · 883d0cc · 883d0cc
1 parent 9e38ec3
commit 883d0cc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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]