Train_Identify/nvidia_ascend_engine/nvidia_engine/MoveEngine/MoveEngine.h

96 lines
2.7 KiB
C++

/*
* @Author: your name
* @Date: 2023-07-21 13:55:06
* @LastEditors: your name
* @LastEditTime: 2023-07-21 14:12:40
* @Description: file content
* @FilePath: /shigetai_lirs/nvidia_ascend_engine/nvidia_engine/MoveEngine/MoveEngine.h
*
* Copyright © 2022 <Shandong Matrix Software Engineering Co., Ltd>
*/
/**
* 动态检测推理engine
* */
#ifndef MOVEENGINE_H
#define MOVEENGINE_H
#include "AppCommon.h"
#include "MyYaml.h"
#include "EngineBase.h"
#include "EngineFactory.h"
#include "retinanet_classify_inference.h"
#include "yolov8_inference.h"
#define STEP0_INPUT_SIZE STEP0_BATCH_SIZE*STEP0_INPUT_CHANNEL*STEP0_INPUT_W*STEP0_INPUT_H
#define STEP0_OUTPUT_ARRAY 4
#define STEP0_OUTPUT_SIZE 1 * STEP0_OUTPUT_ARRAY * 2 // BOX
#define STEP0_CLEAR_NUM 0
#define STEP0_CLASS_NUM 4
class MoveEngine : public ai_matrix::EngineBase
{
public:
MoveEngine();
~MoveEngine();
APP_ERROR Init() override;
APP_ERROR DeInit() override;
APP_ERROR Process() override;
private:
//初始化识别模型
APP_ERROR InitModel();
//获取模型配置
APP_ERROR ReadModelInfo();
//参数初始化
void InitParam();
//使用单device处理
void SingleDeviceProcess(std::shared_ptr<ProcessData> pProcessData, int nType);
void sendComeTrain();
void sendEndTrain();
bool bNeedMoveDetectFlag_;
std::string strPort0_;
ai_matrix::ModelConfig modelConfig_;
std::string strResultPath_;
std::string strBestPath_;
YoloV8Inference yolov8model;
int iStepInter_ = 0; //(0:不识别; 1:开始识别; 2:结束识别)
uint32_t iMoveDataNO_ = 1; //动态检测数据编号
int iHasTrainNum_ = 0; //有车的图象数
int nPreMonitorState = MONITOR_MODEL_INIT_STATE;
std::string strTrainDate_;
std::string strTrainName_;
std::set<int> setPushPort_;
ai_matrix::DataSourceConfig dataSourceCfg_;
std::string INPUT_BLOB_NAME = "images"; // deploy文件中定义的输入层名称
std::string OUTPUT_BLOB_NAME = "output0"; // deploy文件中定义的输出层名称
unsigned int img_width = IMAGE_WIDTH;
unsigned int img_height = IMAGE_HEIGHT;
unsigned int model_width = STEP0_INPUT_W;
unsigned int model_height = STEP0_INPUT_H;
unsigned int clear_num = STEP0_CLEAR_NUM;
unsigned int class_num = STEP0_CLASS_NUM;
unsigned int input_size = STEP0_INPUT_SIZE;
unsigned int output_size = STEP0_OUTPUT_SIZE;
unsigned int det_size = STEP0_CLASS_NUM + STEP0_CLEAR_NUM + 5;
unsigned int batch_size = STEP0_BATCH_SIZE;
float score_threshold = STEP1_SCORE_THRESH;
float nms_threshold = STEP1_NMS_THRESH;
YoloV5ModelInfo modelinfo;
std::queue<std::shared_ptr<ProcessData>> queProcessData_;
};
ENGINE_REGIST(MoveEngine)
#endif