From 41cc7caee64f78b8364db27b5326d1a53b91cb97 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 3 May 2021 15:07:20 +0200
Subject: [PATCH] Add NMS threshold checks (#3020)

* Add NMS threshold checks

* fix
---
 utils/general.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/utils/general.py b/utils/general.py
index 7c06965..fd20a7e 100755
--- 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