#ifndef _YOLOV5_CLASSIFY_INFERENCE_H_
#define _YOLOV5_CLASSIFY_INFERENCE_H_

#include "preprocess.h"
#include "inference.h"
#include "postprocess.h"

using namespace nvinfer1;
using namespace nvcaffeparser1;
using namespace std;

class YoloV5ClassifyInference: public Inference
{
public:
    YoloV5ClassifyInference();
    ~YoloV5ClassifyInference();

    int YoloV5ClassifyInferenceInit(ModelInfo* pYoloV5ClassifyModelInfo, const std::string& strModelName, const std::string& strEngineName);
    int YoloV5ClassifyInferenceDeinit(void);
    int YoloV5ClassifyInferenceModel(cv::Mat& frame, unsigned int* uiClassLabel);

private:
    ModelInfo*  pYoloV5ClassifyModelInfo_ = nullptr;

    std::vector<double> softMax(std::vector<double> vecInput);
};

#endif  //END OF _YOLOV5_CLASSIFY_INFERENCE_H_