/** * 车号属性第一步推理前预处理 * */ #ifndef TRAINSTEPONEENGINE_H #define TRAINSTEPONEENGINE_H #include "AppCommon.h" #include "MyYaml.h" #include "EngineBase.h" #include "EngineFactory.h" #include "yolov5_clearity_inference.h" class TrainStepOneEngine : public ai_matrix::EngineBase { public: TrainStepOneEngine(); ~TrainStepOneEngine(); APP_ERROR Init() override; APP_ERROR DeInit() override; APP_ERROR Process() override; private: //初始化识别模型 APP_ERROR InitModel(); //获取模型配置 APP_ERROR ReadModelInfo(); //设置大框类型 void SetTargetType(PostSubData &postSubData); //过滤无效信息 void FilterInvalidInfo(std::vector &vecRet, std::shared_ptr &pProcessData); //push数据到队列,队列满时则休眠一段时间再push void PushData(const std::string &strPort, const std::shared_ptr &pProcessData); bool bUseEngine_; std::string strPort0_; ai_matrix::ModelConfig modelConfig_; std::string strResultPath_; bool bStepOneImgSaveFlag_; 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 = STEP1_INPUT_W; unsigned int model_height = STEP1_INPUT_H; unsigned int clear_num = STEP1_CLEAR_NUM; unsigned int class_num = STEP1_CLASS_NUM; unsigned int input_size = STEP1_INPUT_SIZE; unsigned int output_size = STEP1_OUTPUT_SIZE; unsigned int det_size = STEP1_CLASS_NUM + STEP1_CLEAR_NUM + 5; unsigned int batch_size = STEP1_BATCH_SIZE; float score_threshold = STEP1_SCORE_THRESH; float nms_threshold = STEP1_NMS_THRESH; YoloV5ClearityModelInfo modelinfo; std::map mapDataSourceCfg_; //[key-数据源id, value-数据源参数配置信息] std::map mapSameSideFlatcarSid_; //[key-数据源id, value-同侧识别平车数据源id] int iSpaceMinRBXPer_; //间隔框最低点不应小于画面某个高度值(该值为画面百分比) }; ENGINE_REGIST(TrainStepOneEngine) #endif