58 lines
1.6 KiB
C++
58 lines
1.6 KiB
C++
//
|
|
// Created by matrixai on 11/9/24.
|
|
//
|
|
|
|
#ifndef TRAIN_CORNERINFERENCEENGINE_H
|
|
#define TRAIN_CORNERINFERENCEENGINE_H
|
|
|
|
#include "AppCommon.h"
|
|
#include "EngineBase.h"
|
|
#include "EngineFactory.h"
|
|
#include "yolov5_clearity_inference.h"
|
|
#include "yolov5_inference.h"
|
|
|
|
using namespace ai_matrix;
|
|
|
|
class CornerInferenceEngine : public ai_matrix::EngineBase
|
|
{
|
|
public:
|
|
CornerInferenceEngine();
|
|
~CornerInferenceEngine();
|
|
|
|
APP_ERROR Init() override;
|
|
APP_ERROR DeInit() override;
|
|
APP_ERROR Process() override;
|
|
|
|
private:
|
|
//初始化识别模型
|
|
APP_ERROR initModel();
|
|
|
|
std::string strPort0_;
|
|
ai_matrix::ModelConfig modelConfig_;
|
|
ai_matrix::DataSourceConfig dataSourceConfig_;
|
|
ai_matrix::IdentifyConfig identifyConfig_;
|
|
// YoloV5ClearityInference yolov5model;
|
|
YoloV5Inference 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 = CORNER_INPUT_W;
|
|
unsigned int model_height = CORNER_INPUT_H;
|
|
unsigned int class_num = CORNER_CLASS_NUM;
|
|
unsigned int input_size = CORNER_INPUT_SIZE;
|
|
unsigned int output_size = CORNER_OUTPUT_SIZE;
|
|
unsigned int det_size = CORNER_CLASS_NUM + 5;
|
|
unsigned int batch_size = CORNER_BATCH_SIZE;
|
|
float score_threshold = CORNER_SCORE_THRESH;
|
|
float nms_threshold = CORNER_NMS_THRESH;
|
|
|
|
// YoloV5ClearityModelInfo modelinfo;
|
|
YoloV5ModelInfo modelinfo;
|
|
// ModelCommonInfo modelCommonInfo;
|
|
};
|
|
ENGINE_REGIST(CornerInferenceEngine)
|
|
|
|
#endif //TRAIN_CORNERINFERENCEENGINE_H
|