following discussion on #11

WongKinYiu · Feb 23, 2024 · 1e33dbb · 1e33dbb
1 parent b2910e1
commit 1e33dbb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/detect.py b/detect.py
@@ -99,7 +99,8 @@ def run(
 
         # NMS
         with dt[2]:
-            pred = non_max_suppression(pred[0][0], conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
+            pred = pred[0][1] if isinstance(pred[0], list) else pred[0]
+            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
 
         # Second-stage classifier (optional)
         # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)