YOLOV5.h 870 Bytes
Newer Older
Your Name's avatar
Your Name committed
1
2
3
4
5
#ifndef __DETECTOR_YOLOV5_H__
#define __DETECTOR_YOLOV5_H__

#include <migraphx/program.hpp>

liucong's avatar
liucong committed
6
#include <CommonDefinition.h>
Your Name's avatar
Your Name committed
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

namespace migraphxSamples
{

typedef struct _YOLOV5Parameter
{
    int numberOfClasses;
    float confidenceThreshold;
    float nmsThreshold;
    float objectThreshold;

}YOLOV5Parameter;

class DetectorYOLOV5
{
public:
    DetectorYOLOV5();
    
    ~DetectorYOLOV5();

    ErrorCode Initialize(InitializationParameterOfDetector initializationParameterOfDetector);

    ErrorCode Detect(const cv::Mat &srcImage, std::vector<ResultOfDetection> &resultsOfDetection);

private:
    cv::FileStorage configurationFile;

    migraphx::program net;
    cv::Size inputSize;
liucong's avatar
liucong committed
36
    std::string inputName;
Your Name's avatar
Your Name committed
37
38
39
    migraphx::shape inputShape;
    
    bool useFP16;
liucong's avatar
liucong committed
40
    std::vector<std::string> classNames;
Your Name's avatar
Your Name committed
41
42
43
44
45
46
47
48
49

    YOLOV5Parameter yolov5Parameter;
};

}


#endif