Add YOLOE detection

README.md

@@ -6,6 +6,7 @@
 
 ## Latest Features
 
+- Add YOLOE detection support [THU-MIG/yoloe](https://github.com/THU-MIG/yoloe)
 - Add new SOTA: YOLOv26, YOLOv26-obb and YOLOv26-seg models from [ultralytics/ultralytics](https://github.com/ultralytics/ultralytics)
 - Add RT-DETRv4 (API similar D-FINE) detection model [RT-DETRs/RT-DETRv4](https://github.com/RT-DETRs/RT-DETRv4)
 - Add D-FINE seg detection model [ArgoHA/D-FINE-seg](https://github.com/ArgoHA/D-FINE-seg)

data/settings_yoloe_seg.ini

@@ -0,0 +1,141 @@
+[detection]
+
+#-----------------------------
+# opencv_dnn = 6
+# tensorrt = 5
+detector_backend = 5
+
+#-----------------------------
+# Target and backend for opencv_dnn detector
+# DNN_TARGET_CPU
+# DNN_TARGET_OPENCL
+# DNN_TARGET_OPENCL_FP16
+# DNN_TARGET_MYRIAD
+# DNN_TARGET_CUDA
+# DNN_TARGET_CUDA_FP16
+ocv_dnn_target = DNN_TARGET_CPU
+
+# DNN_BACKEND_DEFAULT
+# DNN_BACKEND_HALIDE
+# DNN_BACKEND_INFERENCE_ENGINE
+# DNN_BACKEND_OPENCV
+# DNN_BACKEND_VKCOM
+# DNN_BACKEND_CUDA
+# DNN_BACKEND_INFERENCE_ENGINE_NGRAPH
+# DNN_BACKEND_INFERENCE_ENGINE_NN_BUILDER_2019
+ocv_dnn_backend = DNN_BACKEND_OPENCV
+
+#-----------------------------
+nn_weights = C:/work/home/mtracker/Multitarget-tracker/data/yoloe-11m-seg.onnx
+nn_config = C:/work/home/mtracker/Multitarget-tracker/data/yoloe-11m-seg.onnx
+class_names = C:/work/home/mtracker/Multitarget-tracker/data/coco/coco.names
+
+#-----------------------------
+confidence_threshold = 0.3
+
+max_crop_ratio = 0
+max_batch = 1
+gpu_id = 0
+
+#-----------------------------
+# YOLOV3 
+# YOLOV4 
+# YOLOV5 
+net_type = YOLOEMask
+
+#-----------------------------
+# INT8
+# FP16
+# FP32
+inference_precision = FP16
+
+
+[tracking]
+
+#-----------------------------
+# DistCenters = 0   // Euclidean distance between centers, pixels
+# DistRects = 1     // Euclidean distance between bounding rectangles, pixels
+# DistJaccard = 2   // Intersection over Union, IoU, [0, 1]
+# DistHist = 3      // Bhatacharia distance between histograms, [0, 1]
+
+distance_type = 0
+
+#-----------------------------
+# KalmanLinear = 0
+# KalmanUnscented = 1
+
+kalman_type = 0
+
+#-----------------------------
+# FilterCenter = 0
+# FilterRect = 1
+# FilterRRect = 2
+
+filter_goal = 0
+
+#-----------------------------
+# TrackNone = 0
+# TrackKCF = 1
+# TrackMIL = 2
+# TrackMedianFlow = 3
+# TrackGOTURN = 4
+# TrackMOSSE = 5
+# TrackCSRT = 6
+# TrackDAT = 7
+# TrackSTAPLE = 8
+# TrackLDES = 9
+# TrackDaSiamRPN = 10
+# Used if filter_goal == FilterRect
+
+lost_track_type = 0
+
+#-----------------------------
+# MatchHungrian = 0
+# MatchBipart = 1
+
+match_type = 0
+
+#-----------------------------
+# Use constant acceleration motion model:
+# 0 - unused (stable)
+# 1 - use acceleration in Kalman filter (experimental)
+use_aceleration = 0
+
+#-----------------------------
+# Delta time for Kalman filter
+delta_time = 0.4
+
+#-----------------------------
+# Accel noise magnitude for Kalman filter
+accel_noise = 0.2
+
+#-----------------------------
+# Distance threshold between region and object on two frames
+dist_thresh = 0.8 
+
+#-----------------------------
+# If this value > 0 than will be used circle with this radius
+# If this value <= 0 than will be used ellipse with size (3*vx, 3*vy), vx and vy - horizontal and vertical speed in pixelsa
+min_area_radius_pix = -1
+
+#-----------------------------
+# Minimal area radius in ration for object size. Used if min_area_radius_pix < 0
+min_area_radius_k = 0.8
+
+#-----------------------------
+# If the object do not assignment more than this seconds then it will be removed
+max_lost_time = 2
+
+#-----------------------------
+# The maximum trajectory length
+max_trace_len = 2
+
+#-----------------------------
+# Detection abandoned objects
+detect_abandoned = 0
+# After this time (in seconds) the object is considered abandoned
+min_static_time = 5
+# After this time (in seconds) the abandoned object will be removed
+max_static_time = 25
+# Speed in pixels. If speed of object is more that this value than object is non static
+max_speed_for_static = 10

example/VideoExample.cpp

@@ -532,13 +532,19 @@ void VideoExample::DrawTrack(cv::Mat frame,
                              const std::string& userLabel)
 {
     cv::Scalar color = track.m_isStatic ? cv::Scalar(255, 0, 255) : cv::Scalar(0, 255, 0);
+    cv::Rect brect = track.m_rrect.boundingRect();
+
+#if 0
     cv::Point2f rectPoints[4];
     track.m_rrect.points(rectPoints);
     //std::cout << "track: rrect [" << track.m_rrect.size << " from " << track.m_rrect.center << ", " << track.m_rrect.angle << "]" << std::endl;
     for (int i = 0; i < 4; ++i)
     {
         cv::line(frame, rectPoints[i], rectPoints[(i+1) % 4], color);
     }
+#else
+    cv::rectangle(frame, brect, color);
+#endif
 
 #if 0
 #if 0
@@ -607,7 +613,6 @@ void VideoExample::DrawTrack(cv::Mat frame,
         }
     }
 
-    cv::Rect brect = track.m_rrect.boundingRect();
     std::stringstream label;
     label << track.m_ID.ID2Str();
     if (track.m_type != bad_type)

example/examples.h

@@ -543,7 +543,7 @@ class ONNXTensorRTExample final : public VideoExample
 			}
 		}
 
-		//m_detector->CalcMotionMap(frame);
+		m_detector->CalcMotionMap(frame, true);
 	}
 };
 

src/Detector/BaseDetector.h

@@ -159,7 +159,7 @@ class BaseDetector
     /// \brief CalcMotionMap
     /// \param frame
     ///
-    virtual void CalcMotionMap(cv::Mat& frame)
+    virtual void CalcMotionMap(cv::Mat& frame, bool drawOnlyMasks = true)
     {
         if (m_motionMap.size() != frame.size())
             m_motionMap = cv::Mat(frame.size(), CV_32FC1, cv::Scalar(0, 0, 0));
@@ -169,7 +169,8 @@ class BaseDetector
         {
             if (region.m_boxMask.empty())
             {
-                cv::ellipse(foreground, region.m_rrect, cv::Scalar(255, 255, 255), cv::FILLED);
+                if (!drawOnlyMasks)
+                    cv::ellipse(foreground, region.m_rrect, cv::Scalar(255, 255, 255), cv::FILLED);
             }
             else
             {

src/Detector/MotionDetector.cpp

@@ -100,7 +100,7 @@ void MotionDetector::ResetModel(const cv::UMat& img, const cv::Rect& roiRect)
 /// \brief MotionDetector::CalcMotionMap
 /// \param frame
 ///
-void MotionDetector::CalcMotionMap(cv::Mat& frame)
+void MotionDetector::CalcMotionMap(cv::Mat& frame, bool /*drawOnlyMasks*/)
 {
 	if (m_motionMap.size() != frame.size())
 		m_motionMap = cv::Mat(frame.size(), CV_32FC1, cv::Scalar(0, 0, 0));

src/Detector/MotionDetector.h

@@ -22,7 +22,7 @@ class MotionDetector : public BaseDetector
         return true;
     }
 
-    void CalcMotionMap(cv::Mat& frame) override;
+    void CalcMotionMap(cv::Mat& frame, bool drawOnlyMasks) override;
 
     void ResetModel(const cv::UMat& img, const cv::Rect& roiRect) override;
 

src/Detector/OCVDNNDetector.cpp

@@ -176,6 +176,8 @@ bool OCVDNNDetector::Init(const config_t& config)
         dictNetType["YOLOV26"] = ModelType::YOLOV26;
         dictNetType["YOLOV26_OBB"] = ModelType::YOLOV26_OBB;
         dictNetType["YOLOV26Mask"] = ModelType::YOLOV26Mask;
+        dictNetType["YOLOE"] = ModelType::YOLOE;
+        dictNetType["YOLOEMask"] = ModelType::YOLOEMask;
 
         auto netType = dictNetType.find(net_type->second);
         if (netType != dictNetType.end())
@@ -383,6 +385,14 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
     std::vector<cv::Mat> detections;
     m_net.forward(detections, m_outNames); //compute output
 
+#if 0
+    std::cout << "Out layers:\n";
+    for (size_t i = 0; i < detections.size(); ++i)
+    {
+        std::cout << i << ": " << detections[i].size() << "\n";
+    }
+#endif
+
 	switch (m_netType)
 	{
 	case ModelType::YOLOV5:
@@ -448,6 +458,14 @@ void OCVDNNDetector::DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& cr
         ParseYOLOv26_seg(crop, detections, tmpRegions);
         break;
 
+    case ModelType::YOLOE:
+        ParseYOLOE(crop, detections, tmpRegions);
+        break;
+
+    case ModelType::YOLOEMask:
+        ParseYOLOEMask(crop, detections, tmpRegions);
+        break;
+
 	default:
         ParseOldYOLO(crop, detections, tmpRegions);
         break;
@@ -1225,3 +1243,97 @@ void OCVDNNDetector::ParseYOLOv26_seg(const cv::Rect& crop, std::vector<cv::Mat>
         }
     }
 }
+
+///
+/// \brief OCVDNNDetector::ParseYOLOE
+/// \param crop
+/// \param detections
+/// \param tmpRegions
+///
+void OCVDNNDetector::ParseYOLOE(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions)
+{
+    float x_factor = crop.width / static_cast<float>(m_inWidth);
+    float y_factor = crop.height / static_cast<float>(m_inHeight);
+
+    //0: name: images, size: 1x3x640x640
+    //1: name: output0, size: 1x54x8400
+    //2: name: output1, size: 1x32x160x160
+
+    int rows = detections[0].size[2];
+    int dimensions = detections[0].size[1];
+
+	detections[0] = detections[0].reshape(1, dimensions);
+	cv::transpose(detections[0], detections[0]);
+
+    float* data = (float*)detections[0].data;
+
+    for (int i = 0; i < rows; ++i)
+    {
+        float* classes_scores = data + 4;
+
+        cv::Mat scores(1, static_cast<int>(m_classNames.size()), CV_32FC1, classes_scores);
+        cv::Point class_id;
+        double maxClassScore = 0;
+        cv::minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
+
+        if (maxClassScore > m_confidenceThreshold)
+        {
+            float x = data[0] * x_factor + crop.x;
+            float y = data[1] * y_factor + crop.y;
+            float w = data[2] * x_factor;
+            float h = data[3] * y_factor;
+            //float angle = 180.f * data[4 + scores.cols] / M_PI;
+
+            if (m_classesWhiteList.empty() || m_classesWhiteList.find(T2T(class_id.x)) != std::end(m_classesWhiteList))
+                tmpRegions.emplace_back(cv::RotatedRect(cv::Point2f(x, y), cv::Size2f(w, h), 0), T2T(class_id.x), static_cast<float>(maxClassScore));
+        }
+        data += dimensions;
+    }
+}
+
+///
+/// \brief OCVDNNDetector::ParseYOLOEMask
+/// \param crop
+/// \param detections
+/// \param tmpRegions
+///
+void OCVDNNDetector::ParseYOLOEMask(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions)
+{
+    float x_factor = crop.width / static_cast<float>(m_inWidth);
+    float y_factor = crop.height / static_cast<float>(m_inHeight);
+
+    //0: name: images, size: 1x3x640x640
+    //1: name: output0, size: 1x54x8400
+    //2: name: output1, size: 1x32x160x160
+
+    int rows = detections[0].size[2];
+    int dimensions = detections[0].size[1];
+
+    detections[0] = detections[0].reshape(1, dimensions);
+    cv::transpose(detections[0], detections[0]);
+
+    float* data = (float*)detections[0].data;
+
+    for (int i = 0; i < rows; ++i)
+    {
+        float* classes_scores = data + 4;
+
+        cv::Mat scores(1, static_cast<int>(m_classNames.size()), CV_32FC1, classes_scores);
+        cv::Point class_id;
+        double maxClassScore = 0;
+        cv::minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
+
+        if (maxClassScore > m_confidenceThreshold)
+        {
+            float x = data[0] * x_factor + crop.x;
+            float y = data[1] * y_factor + crop.y;
+            float w = data[2] * x_factor;
+            float h = data[3] * y_factor;
+            //float angle = 180.f * data[4 + scores.cols] / M_PI;
+
+            if (m_classesWhiteList.empty() || m_classesWhiteList.find(T2T(class_id.x)) != std::end(m_classesWhiteList))
+                tmpRegions.emplace_back(cv::RotatedRect(cv::Point2f(x, y), cv::Size2f(w, h), 0), T2T(class_id.x), static_cast<float>(maxClassScore));
+        }
+        data += dimensions;
+    }
+}

src/Detector/OCVDNNDetector.h

@@ -56,7 +56,9 @@ class OCVDNNDetector final : public BaseDetector
         DFINE_IS,
         YOLOV26,
         YOLOV26_OBB,
-        YOLOV26Mask
+        YOLOV26Mask,
+        YOLOE,
+        YOLOEMask
     };
 
     cv::dnn::Net m_net;
@@ -97,6 +99,8 @@ class OCVDNNDetector final : public BaseDetector
     void ParseYOLOv26(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
     void ParseYOLOv26_obb(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
     void ParseYOLOv26_seg(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
+    void ParseYOLOE(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
+    void ParseYOLOEMask(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
 };
 
 #endif

src/Detector/ONNXTensorRTDetector.cpp

@@ -106,6 +106,8 @@ bool ONNXTensorRTDetector::Init(const config_t& config)
 		dictNetType["YOLOV26"] = tensor_rt::YOLOV26;
 		dictNetType["YOLOV26_OBB"] = tensor_rt::YOLOV26_OBB;
 		dictNetType["YOLOV26Mask"] = tensor_rt::YOLOV26Mask;
+		dictNetType["YOLOE"] = tensor_rt::YOLOE;
+		dictNetType["YOLOEMask"] = tensor_rt::YOLOEMask;
 
 		auto netType = dictNetType.find(net_type->second);
 		if (netType != dictNetType.end())