generated from zhangwei/Train_Identify
69 lines
2.2 KiB
C++
69 lines
2.2 KiB
C++
/**
|
||
* 车号属性第一步推理前预处理
|
||
* */
|
||
|
||
#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<stDetection> &vecRet, std::shared_ptr<ProcessData> &pProcessData);
|
||
//push数据到队列，队列满时则休眠一段时间再push
|
||
void PushData(const std::string &strPort, const std::shared_ptr<ProcessData> &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<int, ai_matrix::DataSourceConfig> mapDataSourceCfg_; //[key-数据源id, value-数据源参数配置信息]
|
||
std::map<int, int> mapSameSideFlatcarSid_; //[key-数据源id, value-同侧识别平车数据源id]
|
||
int iSpaceMinRBXPer_; //间隔框最低点不应小于画面某个高度值(该值为画面百分比)
|
||
|
||
};
|
||
|
||
ENGINE_REGIST(TrainStepOneEngine)
|
||
|
||
#endif
|