From 917d3453754bcd9402fce7ef50fa27dd672b95a4 Mon Sep 17 00:00:00 2001
From: Marcos Luciano <marcoslucianops@gmail.com>
Date: Sun, 27 Dec 2020 12:31:54 -0300
Subject: [PATCH] Updated YOLOv5

---
 YOLOv5.md                                          | 12 +++++++++---
 external/yolov5/config_infer_primary.txt           |  3 +++
 .../nvdsparsebbox_Yolo.cpp                         | 14 +++-----------
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/YOLOv5.md b/YOLOv5.md
index dcb8dad..f872fe5 100644
--- a/YOLOv5.md
+++ b/YOLOv5.md
@@ -198,9 +198,15 @@ to
 model-engine-file=yolov5x.engine
 ```
 
-To change NMS_THRESH and THRESH_CONF, edit nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp file and recompile
+To change NMS_THRESH, edit nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp file and recompile
 
 ```
-#define NMS_THRESH 0.45
-#define CONF_THRESH 0.25
+#define kNMS_THRESH 0.45
+```
+
+To change CONF_THRESH, edit config_infer_primary.txt file
+
+```
+[class-attrs-all]
+pre-cluster-threshold=0.25
 ```
diff --git a/external/yolov5/config_infer_primary.txt b/external/yolov5/config_infer_primary.txt
index 061399e..a28a5a8 100644
--- a/external/yolov5/config_infer_primary.txt
+++ b/external/yolov5/config_infer_primary.txt
@@ -14,3 +14,6 @@ maintain-aspect-ratio=0
 parse-bbox-func-name=NvDsInferParseCustomYoloV5
 custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so
 engine-create-func-name=NvDsInferYoloCudaEngineGet
+
+[class-attrs-all]
+pre-cluster-threshold=0.25
diff --git a/external/yolov5/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp b/external/yolov5/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp
index 55ee116..e38d4c1 100644
--- a/external/yolov5/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp
+++ b/external/yolov5/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp
@@ -25,10 +25,7 @@
 
 #include <map>
 
-#define NMS_THRESH 0.45
-#define CONF_THRESH 0.25
-
-static const int NUM_CLASSES_YOLO = 80;
+#define kNMS_THRESH 0.45
 
 static constexpr int LOCATIONS = 4;
 struct alignas(float) Detection{
@@ -90,16 +87,11 @@ static bool NvDsInferParseYoloV5(
     NvDsInferParseDetectionParams const& detectionParams,
     std::vector<NvDsInferParseObjectInfo>& objectList)
 {
-    if (NUM_CLASSES_YOLO != detectionParams.numClassesConfigured)
-    {
-        std::cerr << "WARNING: Num classes mismatch. Configured:"
-                  << detectionParams.numClassesConfigured
-                  << ", detected by network: " << NUM_CLASSES_YOLO << std::endl;
-    }
+    const float kCONF_THRESH = detectionParams.perClassThreshold[0];
 
     std::vector<Detection> res;
 
-    nms(res, (float*)(outputLayersInfo[0].buffer), CONF_THRESH, NMS_THRESH);
+    nms(res, (float*)(outputLayersInfo[0].buffer), kCONF_THRESH, kNMS_THRESH);
     
     for(auto& r : res) {
 	    NvDsInferParseObjectInfo oinfo;