YOLOv10 with opencv_dnn

Smorodov · Jun 9, 2024 · 82c2e30 · 82c2e30
1 parent b649f2d
commit 82c2e30
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/src/Detector/OCVDNNDetector.cpp b/src/Detector/OCVDNNDetector.cpp
@@ -369,7 +369,7 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
 
             for (int i = 0; i < rows; ++i)
             {
-                if (m_netType == ModelType::YOLOV8 || m_netType == ModelType::YOLOV9 || m_netType == ModelType::YOLOV10)
+                if (m_netType == ModelType::YOLOV8 || m_netType == ModelType::YOLOV9)
                 {
                     float* classes_scores = data + 4;
 
@@ -395,7 +395,7 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
                             tmpRegions.emplace_back(cv::Rect(left + crop.x, top + crop.y, width, height), T2T(class_id.x), static_cast<float>(maxClassScore));
                     }
                 }
-                else // yolov5
+                else if (m_netType == ModelType::YOLOV5)
                 {
                     float confidence = data[4];
 
@@ -426,6 +426,21 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
                         }
                     }
                 }
+                else if (m_netType == ModelType::YOLOV10)
+                {
+                    int left = cvRound(x_factor * data[0]);
+                    int top = cvRound(y_factor * data[1]);
+                    int width = cvRound(x_factor * (data[2] - data[0]));
+                    int height = cvRound(y_factor * (data[3] - data[1]));
+                    float confidence = data[4];
+                    int classId = cvRound(data[5]);
+
+                    if (confidence >= m_confidenceThreshold)
+                    {
+                        if (m_classesWhiteList.empty() || m_classesWhiteList.find(T2T(classId)) != std::end(m_classesWhiteList))
+                            tmpRegions.emplace_back(cv::Rect(left + crop.x, top + crop.y, width, height), T2T(classId), confidence);
+                    }
+                }
 
                 data += dimensions;
             }