Multitarget-tracker/src/Detector/OCVDNNDetector.h at cursor/create-project-tutorial-d481 · helloworld1301/Multitarget-tracker · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#pragma once

#ifdef USE_OCV_DNN

#include "BaseDetector.h"

#include <opencv2/dnn.hpp>
#include <opencv2/dnn/shape_utils.hpp>

///
/// \brief The OCVDNNDetector class
///
class OCVDNNDetector final : public BaseDetector
{
public:
    OCVDNNDetector(const cv::UMat& colorFrame);
    OCVDNNDetector(const cv::Mat& colorFrame);
    ~OCVDNNDetector(void) = default;

    bool Init(const config_t& config) override;

    void Detect(const cv::UMat& colorFrame) override;

    bool CanGrayProcessing() const override
    {
        return false;
    }

private:
    enum class ModelType
    {
        Unknown,
        YOLOV3,
        YOLOV3_TINY,
        YOLOV4,
        YOLOV4_TINY,
        YOLOV5,
        YOLOV5_OBB,
        YOLOV5Mask,
        YOLOV6,
        YOLOV7,
        YOLOV7Mask,
        YOLOV8,
        YOLOV8_OBB,
        YOLOV8Mask,
        YOLOV9,
        YOLOV10,
        YOLOV11,
        YOLOV11_OBB,
        YOLOV11Mask,
        YOLOV12,
        RFDETR
    };

    cv::dnn::Net m_net;

    void DetectInCrop(const cv::UMat& colorFrame, const cv::Rect& crop, regions_t& tmpRegions);

    int m_inWidth = 608;
    int m_inHeight = 608;

    float m_WHRatio = 1.f;
    float m_inScaleFactor = 0.003921f;
    float m_meanVal = 0.f;
    float m_confidenceThreshold = 0.24f;
    track_t m_nmsThreshold = static_cast<track_t>(0.4);
    bool m_swapRB = false;
    float m_maxCropRatio = 2.0f;
    ModelType m_netType = ModelType::Unknown;
    std::vector<std::string> m_classNames;
    std::vector<cv::String> m_outNames;
    std::vector<int> m_outLayers;
    std::string m_outLayerType;
    cv::UMat m_inputBlob;

    void ParseOldYOLO(const cv::Rect& crop, const std::vector<cv::Mat>& detections, regions_t& tmpRegions);

    void ParseYOLOv5(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv8(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv9(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv10(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv11(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv5_8_11_obb(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseYOLOv5_8_11_seg(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
    void ParseRFDETR(const cv::Rect& crop, std::vector<cv::Mat>& detections, regions_t& tmpRegions);
};

#endif