From a4724d55924708b4cee0dacbc875545b4b3abe3f Mon Sep 17 00:00:00 2001
From: Mohammad Khoshbin <khoshbin.mohammad.mk@gmail.com>
Date: Wed, 27 Jul 2022 19:17:28 +0430
Subject: [PATCH] non-max-suppression function fixed for one class model conf.
 (#305)

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

diff --git a/utils/general.py b/utils/general.py
index 7dc53eb8..faf908f9 100644
--- a/utils/general.py
+++ b/utils/general.py
@@ -645,7 +645,11 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
             continue
 
         # Compute conf
-        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
+        if nc == 1:
+            x[:, 5:] = x[:, 4:5] # for models with one class, cls_loss is 0 and cls_conf is always 0.5,
+                                 # so there is no need to multiplicate.
+        else:
+            x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
 
         # Box (center x, center y, width, height) to (x1, y1, x2, y2)
         box = xywh2xyxy(x[:, :4])