Train_Identify_arm/nvidia_ascend_engine/nvidia_engine/ChkDateStepOneEngine/ChkDateStepOneEngine.h

63 lines
1.8 KiB
C
Raw Permalink Normal View History

2024-06-19 06:35:05 +00:00
/**
*
* */
#ifndef CHKDATESTEPONEENGINE_H
#define CHKDATESTEPONEENGINE_H
#include "AppCommon.h"
#include "MyYaml.h"
#include "EngineBase.h"
#include "EngineFactory.h"
#include "yolov5_clearity_inference.h"
class ChkDateStepOneEngine : public ai_matrix::EngineBase
{
public:
ChkDateStepOneEngine();
~ChkDateStepOneEngine();
APP_ERROR Init() override;
APP_ERROR DeInit() override;
APP_ERROR Process() override;
private:
//初始化识别模型
APP_ERROR InitModel();
//获取模型配置
APP_ERROR ReadModelInfo();
//过滤无效信息
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_DATE;
unsigned int input_size = STEP1_INPUT_SIZE;
unsigned int output_size = STEP1_OUTPUT_SIZE;
unsigned int det_size = STEP1_CLASS_NUM_DATE + 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;
};
ENGINE_REGIST(ChkDateStepOneEngine)
#endif