From 6d3d521368bfb589396004f6e7988becc57e3fd9 Mon Sep 17 00:00:00 2001
From: Bubbliiiing <47347516+bubbliiiing@users.noreply.github.com>
Date: Fri, 11 Sep 2020 14:09:15 +0800
Subject: [PATCH] Update yolo.py

---
 yolo.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/yolo.py b/yolo.py
index d9e9226..e060c98 100644
--- a/yolo.py
+++ b/yolo.py
@@ -24,6 +24,7 @@ class YOLO(object):
         "classes_path": 'model_data/coco_classes.txt',
         "model_image_size" : (416, 416, 3),
         "confidence": 0.5,
+        "iou" : 0.3,
         "cuda": True
     }
 
@@ -123,7 +124,7 @@ class YOLO(object):
         output = torch.cat(output_list, 1)
         batch_detections = non_max_suppression(output, len(self.class_names),
                                                 conf_thres=self.confidence,
-                                                nms_thres=0.3)
+                                                nms_thres=self.iou)
         try:
             batch_detections = batch_detections[0].cpu().numpy()
         except:
-- 
GitLab