61 lines
1.7 KiB
C
61 lines
1.7 KiB
C
|
/**
|
|||
|
|
* 车号属性第一步推理前预处理
|
|||
|
|
* */
|
|||
|
|
|
|||
|
|
#ifndef TRAINSTEPTWOENGINE_H
|
|||
|
|
#define TRAINSTEPTWOENGINE_H
|
|||
|
|
|
|||
|
|
#include "AppCommon.h"
|
|||
|
|
#include "MyYaml.h"
|
|||
|
|
#include "EngineBase.h"
|
|||
|
|
#include "EngineFactory.h"
|
|||
|
|
#include "yolov5_clearity_inference.h"
|
|||
|
|
|
|||
|
|
class TrainStepTwoEngine : public ai_matrix::EngineBase
|
|||
|
|
{
|
|||
|
|
public:
|
|||
|
|
TrainStepTwoEngine();
|
|||
|
|
~TrainStepTwoEngine();
|
|||
|
|
|
|||
|
|
APP_ERROR Init() override;
|
|||
|
|
APP_ERROR DeInit() override;
|
|||
|
|
APP_ERROR Process() override;
|
|||
|
|
|
|||
|
|
private:
|
|||
|
|
//初始化识别模型
|
|||
|
|
APP_ERROR InitModel();
|
|||
|
|
//获取模型配置
|
|||
|
|
APP_ERROR ReadModelInfo();
|
|||
|
|
//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 = 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(TrainStepTwoEngine)
|
|||
|
|
|
|||
|
|
#endif
|