Update YOLOv5Detector.cpp

Author: shaoshengsong

detector/YOLOv5/src/YOLOv5Detector.cpp

@@ -4,29 +4,41 @@
     @Date        : 2022-09-23 02:52:41
 */
 #include <YOLOv5Detector.h>
+#include <iostream>
 
-void YOLOv5Detector::init(std::string onnxpath) {
+void YOLOv5Detector::init(std::string onnxpath)
+{
 
     this->net = cv::dnn::readNetFromONNX(onnxpath);
+
+    std::string file="./coco_80_labels_list.txt";
+    std::ifstream ifs(file);
+    if (!ifs.is_open())
+        CV_Error(cv::Error::StsError, "File " + file + " not found");
+    std::string line;
+    while (std::getline(ifs, line))
+    {
+        classes_.push_back(line);
+    }
 }
 
 void YOLOv5Detector::detect(cv::Mat & frame, std::vector<detect_result> &results)
 {
 
-
-
     int w = frame.cols;
     int h = frame.rows;
     int _max = std::max(h, w);
     cv::Mat image = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3);
     cv::Rect roi(0, 0, w, h);
     frame.copyTo(image(roi));
 
-    float x_factor = image.cols / model_input_width_;
-    float y_factor = image.rows / model_input_height_;
-    //std::cout<<image.cols<<":"<<image.rows<<std::endl;
+
+    float x_factor = static_cast<float>(image.cols) / model_input_width_;
+    float y_factor = static_cast<float>(image.rows) / model_input_height_;
 
 
+    std::cout<<x_factor<<std::endl;
+    std::cout<<y_factor<<std::endl;
 
     cv::Mat blob = cv::dnn::blobFromImage(image, 1 / 255.0, cv::Size(model_input_width_, model_input_height_), cv::Scalar(0, 0, 0), true, false);
     this->net.setInput(blob);
@@ -40,18 +52,19 @@ void YOLOv5Detector::detect(cv::Mat & frame, std::vector<detect_result> &results
     std::vector<float> confidences;
     for (int i = 0; i < det_output.rows; i++)
     {
-        float tmp_confidence = det_output.at<float>(i, 4);
-        if (tmp_confidence < nms_threshold_)
+        float box_conf = det_output.at<float>(i, 4);
+        if (box_conf < nms_threshold_)
         {
             continue;
         }
+
         cv::Mat classes_confidences = det_output.row(i).colRange(5, 85);
         cv::Point classIdPoint;
-        double confidence;
-        minMaxLoc(classes_confidences, 0, &confidence, 0, &classIdPoint);
+        double cls_conf;
+        cv::minMaxLoc(classes_confidences, 0, &cls_conf, 0, &classIdPoint);
 
 
-        if (confidence > confidence_threshold_)
+        if (cls_conf > confidence_threshold_)
         {
             float cx = det_output.at<float>(i, 0);
             float cy = det_output.at<float>(i, 1);
@@ -69,7 +82,7 @@ void YOLOv5Detector::detect(cv::Mat & frame, std::vector<detect_result> &results
 
             boxes.push_back(box);
             classIds.push_back(classIdPoint.x);
-            confidences.push_back(confidence);
+            confidences.push_back(cls_conf * box_conf);
         }
     }
 
@@ -83,11 +96,23 @@ void YOLOv5Detector::detect(cv::Mat & frame, std::vector<detect_result> &results
         dr.box = boxes[index];
         dr.classId = idx;
         dr.confidence = confidences[index];
-        cv::rectangle(frame, boxes[index], cv::Scalar(0, 0, 255), 2, 8);
-
-        cv::rectangle(frame, cv::Point(boxes[index].tl().x, boxes[index].tl().y - 20),
-                      cv::Point(boxes[index].br().x, boxes[index].tl().y), cv::Scalar(255, 0, 0), -1);
         results.push_back(dr);
     }
 
 }
+void YOLOv5Detector::draw_frame(cv::Mat & frame, std::vector<detect_result> &results)
+{
+
+    for(auto dr : results)
+    {
+
+        cv::rectangle(frame, dr.box , cv::Scalar(0, 0, 255), 2, 8);
+        cv::rectangle(frame, cv::Point(dr.box .tl().x, dr.box .tl().y - 20), cv::Point(dr.box .br().x, dr.box .tl().y), cv::Scalar(255, 0, 0), -1);
+
+        std::string label = cv::format("%.2f", dr.confidence);
+        label = classes_[dr.classId ] + ":" + label;
+
+        cv::putText(frame, label, cv::Point(dr.box.x, dr.box.y + 6), 1, 2, cv::Scalar(0, 255, 0),2);
+
+    }
+}