/**
 * 集装箱第二步推理前预处理
 * */

#ifndef ContainerStep2InferenceEngine_H
#define ContainerStep2InferenceEngine_H

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

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

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

private:
	//初始化识别模型
	APP_ERROR initModel();
    void resetLocation(SingleData &singleData, SingleData &step1SingleData, float fResizeRatio = 1.0);

    std::string strPort0_;
    ai_matrix::ModelConfig modelConfig_;
    ai_matrix::BaseConfig baseConfig_;
    ai_matrix::IdentifyConfig identifyConfig_;

    YoloV5ClearityInference 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 = STEP2_INPUT_W;
    unsigned int model_height = STEP2_INPUT_H;
    unsigned int clear_num = STEP2_CLEAR_NUM;
    unsigned int class_num = STEP2_CLASS_NUM;
    unsigned int input_size = STEP2_INPUT_SIZE;
    unsigned int output_size = STEP2_OUTPUT_SIZE;
    unsigned int det_size = STEP2_CLASS_NUM + STEP2_CLEAR_NUM + 5;
    unsigned int batch_size = STEP2_BATCH_SIZE;
    float score_threshold = STEP2_SCORE_THRESH;
    float nms_threshold = STEP2_NMS_THRESH;

	YoloV5ClearityModelInfo modelinfo;
	
};

ENGINE_REGIST(ContainerStep2InferenceEngine)

#endif