YOLOV5.h

#ifndef __DETECTOR_YOLOV5_H__
#define __DETECTOR_YOLOV5_H__

#include <migraphx/program.hpp>
#include <CommonDefinition.h>

namespace migraphxSamples
{

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

class DetectorYOLOV5
{
    public:
    DetectorYOLOV5();

    ~DetectorYOLOV5();

    ErrorCode Initialize(InitializationParameterOfDetector initializationParameterOfDetector,
                         bool dynamic);

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

    private:
    cv::FileStorage configurationFile;

    migraphx::program net;
    cv::Size inputSize;
    std::string inputName;
    std::string modelPath;
    migraphx::shape inputShape;

    bool useFP16;
    std::vector<std::string> classNames;

    YOLOV5Parameter yolov5Parameter;
};

} // namespace migraphxSamples

#endif