Add NMS threshold checks (#3020)

* Add NMS threshold checks * fix
ultralytics · May 3, 2021 · 41cc7ca · 41cc7ca
1 parent fcd5702
commit 41cc7ca
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/utils/general.py b/utils/general.py
@@ -476,6 +476,10 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
     nc = prediction.shape[2] - 5  # number of classes
     xc = prediction[..., 4] > conf_thres  # candidates
 
+    # Checks
+    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
+    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
+
     # Settings
     min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
     max_det = 300  # maximum number of detections per image