From 57985c8d32e6d81a0ba73cf2a99e22c0f0341859 Mon Sep 17 00:00:00 2001
From: Ninenox <nisit15@hotmail.com>
Date: Mon, 21 Aug 2023 21:23:12 +0700
Subject: [PATCH] add nms for fixed stack bouding boxes object detection

---
 models/torch_utils.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/models/torch_utils.py b/models/torch_utils.py
index 36d0255..f4e80ab 100644
--- a/models/torch_utils.py
+++ b/models/torch_utils.py
@@ -61,6 +61,7 @@ def pose_postprocess(
 
 def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]):
     assert len(data) == 4
+    iou_thres: float = 0.65
     num_dets, bboxes, scores, labels = data[0][0], data[1][0], data[2][
         0], data[3][0]
     nums = num_dets.item()
@@ -70,6 +71,8 @@ def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]):
     bboxes = bboxes[:nums]
     scores = scores[:nums]
     labels = labels[:nums]
+    idx = nms(bboxes, scores, iou_thres)
+    bboxes, scores, labels = bboxes[idx], scores[idx], labels[idx]
     return bboxes, scores, labels