//
// Created by matrixai on 11/9/24.
//

#ifndef TRAIN_CORNERINFERENCEENGINE_H
#define TRAIN_CORNERINFERENCEENGINE_H

#include "AppCommon.h"
#include "EngineBase.h"
#include "EngineFactory.h"
#include "yolov5_clearity_inference.h"
#include "yolov5_inference.h"

using namespace ai_matrix;

class CornerInferenceEngine  : public ai_matrix::EngineBase
{
public:
    CornerInferenceEngine();
    ~CornerInferenceEngine();

    APP_ERROR Init() override;
    APP_ERROR DeInit() override;
    APP_ERROR Process() override;

private:
    //初始化识别模型
    APP_ERROR initModel();

    std::string strPort0_;
    ai_matrix::ModelConfig modelConfig_;
    ai_matrix::DataSourceConfig dataSourceConfig_;
    ai_matrix::IdentifyConfig identifyConfig_;
//    YoloV5ClearityInference yolov5model;
    YoloV5Inference yolov5model;

    const char*  INPUT_BLOB_NAME = "images";   //输入层名称
    const char*  OUTPUT_BLOB_NAME = "output";   //输出层名称
    unsigned int img_width = IMAGE_WIDTH;
    unsigned int img_height = IMAGE_HEIGHT;
    unsigned int model_width = CORNER_INPUT_W;
    unsigned int model_height = CORNER_INPUT_H;
    unsigned int class_num = CORNER_CLASS_NUM;
    unsigned int input_size = CORNER_INPUT_SIZE;
    unsigned int output_size = CORNER_OUTPUT_SIZE;
    unsigned int det_size = CORNER_CLASS_NUM + 5;
    unsigned int batch_size = CORNER_BATCH_SIZE;
    float score_threshold = CORNER_SCORE_THRESH;
    float nms_threshold = CORNER_NMS_THRESH;

//    YoloV5ClearityModelInfo modelinfo;
    YoloV5ModelInfo modelinfo;
//    ModelCommonInfo modelCommonInfo;
};
ENGINE_REGIST(CornerInferenceEngine)

#endif //TRAIN_CORNERINFERENCEENGINE_H