From 8ca0568017134f6494da28c336a89bb3b4a9a470 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20Krau=C3=9Fer?= <koro.no.omo@gmail.com>
Date: Sun, 31 May 2026 22:02:30 +0200
Subject: [PATCH] added non max supression

---
 api/src/yolo/train_yolo_faces.py      |  2 +-
 api/src/yolo/yolo_model_production.py | 34 ++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/api/src/yolo/train_yolo_faces.py b/api/src/yolo/train_yolo_faces.py
index 49b9702..e3b20a9 100644
--- a/api/src/yolo/train_yolo_faces.py
+++ b/api/src/yolo/train_yolo_faces.py
@@ -75,7 +75,7 @@ def use_webcam(grid, img_size):
 
         prediction = model(image)
 
-        bboxes, grid_ob, grid_noob = convert_prediction(prediction.squeeze(0), image.squeeze(0), threshold=0.95)
+        bboxes, grid_ob, grid_noob = convert_prediction(prediction.squeeze(0), image.squeeze(0), threshold=0.9)
 
         for bbox in grid_noob:
             xmin = int(bbox[0] * scale_w)
diff --git a/api/src/yolo/yolo_model_production.py b/api/src/yolo/yolo_model_production.py
index 920a813..82d6531 100644
--- a/api/src/yolo/yolo_model_production.py
+++ b/api/src/yolo/yolo_model_production.py
@@ -29,7 +29,8 @@ def convert_prediction(label, image, threshold=0.9):
 
             boxx, boxy = turn_image_centered(x=boxx, y=boxy, img_w=image_size, img_h=image_size, S=grid_number, cell_i=x, cell_j=y)
 
-            boxes_to_draw.append(xy_center_to_edges(boxx, boxy, boxw, boxh)) #xmin, ymin, xmax, ymax
+            boxes_to_draw.append([label[x, y, 4]] + xy_center_to_edges(boxx, boxy, boxw, boxh)) #xmin, ymin, xmax, ymax
+    boxes_to_draw = nms(boxes_to_draw)
     return boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj
 
 def turn_image_centered(x, y, img_w, img_h, S, cell_i, cell_j):
@@ -50,6 +51,37 @@ def xy_center_to_edges(xcenter, ycenter, width, height):
 
     return [x, y, x + width, y + height]
 
+def iou(boxA, boxB):
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+    inter_area = max(0, xB - xA) * max(0, yB - yA)
+
+    boxA_area = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
+    boxB_area = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
+
+    union = boxA_area + boxB_area - inter_area
+
+    return inter_area / union if union > 0 else 0
+
+def nms(bboxes, iou_threshold=0.1):
+    bboxes = sorted(bboxes, key=lambda x: x[0], reverse=True)
+
+    keep = []
+
+    while bboxes:
+        best = bboxes.pop(0)
+        keep.append(best[1:5])
+
+        bboxes = [
+            box for box in bboxes
+            if iou(best[1:5], box[1:5]) < iou_threshold
+        ]
+
+    return keep
+
 class Yolo_Conv_Block(nn.Module):
     def __init__(self, c_in, c_hidden, c_out, kernel_size):
         super().__init__()