185 lines
7.7 KiB
C++
185 lines
7.7 KiB
C++
#ifndef _MODEL_H_
|
|
#define _MODEL_H_
|
|
|
|
#include <iostream>
|
|
|
|
using namespace std;
|
|
|
|
#define USE_FP16 // set USE_INT8 or USE_FP16 or USE_FP32
|
|
#define DEVICE 0 // GPU id
|
|
#define BATCH_SIZE 1 //batch size
|
|
#define MAX_WORKSPAXE_SIZE (16 * (1 << 20)) //工作空间大小
|
|
#define MAX_IMAGE_INPUT_SIZE_THRESH 3000*3000
|
|
#define INPUT_CHANNEL 3 //输入通道数
|
|
#define LOCATIONS 4
|
|
|
|
//yolov5 目标检测
|
|
|
|
// == step 1 👇 yolov5 ==
|
|
// -- 车号/属性 --
|
|
#define STEP1_INPUT_BLOB_NAME "images" //输入层名称
|
|
#define STEP1_OUTPUT_BLOB_NAME "output" //输出层名称
|
|
|
|
#define STEP1_INPUT_CHANNEL 3 //输入通道数
|
|
#define STEP1_BATCH_SIZE 1 //batch size
|
|
|
|
#define STEP1_POS_CONF 5 //置信度加坐标
|
|
#define STEP1_LOCATIONS 4 //坐标数目
|
|
|
|
#define STEP1_NMS_THRESH 0.5 //step1 非极大值抑制阈值
|
|
#define STEP1_SCORE_THRESH 0.9 //step1 置信度(得分)阈值
|
|
|
|
#define STEP1_CLASS_NUM 19 //step1 分类数
|
|
#define STEP1_CLEAR_NUM 21 //step1 清晰度个数
|
|
|
|
#define STEP1_INPUT_H 960 //step1 输入图像高
|
|
#define STEP1_INPUT_W 960 //step1 输入图像宽
|
|
#define STEP1_BATCH_SIZE 1
|
|
#define STEP1_OUTPUT_HISTOGRAM_N 5 //不能超过139167
|
|
|
|
#define STEP1_BBOX_SIZE1 (STEP1_INPUT_H/32)
|
|
#define STEP1_BBOX_SIZE2 (STEP1_INPUT_H/16)
|
|
#define STEP1_BBOX_SIZE3 (STEP1_INPUT_H/8)
|
|
|
|
#define STEP1_DET_SIZE (STEP1_CLASS_NUM + STEP1_CLEAR_NUM + STEP1_POS_CONF)
|
|
|
|
#define STEP1_INPUT_SIZE INPUT_CHANNEL*STEP1_INPUT_H*STEP1_INPUT_W //input
|
|
#define STEP1_OUTPUT_SIZE INPUT_CHANNEL*(STEP1_BBOX_SIZE1*STEP1_BBOX_SIZE1+STEP1_BBOX_SIZE2*STEP1_BBOX_SIZE2+STEP1_BBOX_SIZE3*STEP1_BBOX_SIZE3)*(STEP1_CLEAR_NUM+STEP1_CLASS_NUM+STEP1_POS_CONF) //output
|
|
// -- 集装箱 --
|
|
#define STEP1_CONTAINER_INPUT_BLOB_NAME "images" //输入层名称
|
|
#define STEP1_CONTAINER_OUTPUT_BLOB_NAME "output" //输出层名称
|
|
#define STEP1_CONTAINER_INPUT_CHANNEL 3 //输入通道数
|
|
#define STEP1_CONTAINER_BATCH_SIZE 1 //batch size
|
|
#define STEP1_CONTAINER_POS_CONF 5 //置信度加坐标
|
|
#define STEP1_CONTAINER_LOCATIONS 4 //坐标数目
|
|
#define STEP1_CONTAINER_NMS_THRESH 0.5 //step1 非极大值抑制阈值
|
|
#define STEP1_CONTAINER_SCORE_THRESH 0.9 //step1 置信度(得分)阈值
|
|
#define STEP1_CONTAINER_CLASS_NUM 19 //step1 分类数
|
|
#define STEP1_CONTAINER_CLEAR_NUM 21 //step1 清晰度个数
|
|
#define STEP1_CONTAINER_INPUT_H 960 //step1 输入图像高
|
|
#define STEP1_CONTAINER_INPUT_W 960 //step1 输入图像宽
|
|
#define STEP1_CONTAINER_BATCH_SIZE 1
|
|
#define STEP1_CONTAINER_BBOX_SIZE1 (STEP1_CONTAINER_INPUT_H/32)
|
|
#define STEP1_CONTAINER_BBOX_SIZE2 (STEP1_CONTAINER_INPUT_H/16)
|
|
#define STEP1_CONTAINER_BBOX_SIZE3 (STEP1_CONTAINER_INPUT_H/8)
|
|
#define STEP1_CONTAINER_DET_SIZE (STEP1_CONTAINER_CLASS_NUM + STEP1_CONTAINER_CLEAR_NUM + STEP1_CONTAINER_POS_CONF)
|
|
#define STEP1_CONTAINER_INPUT_SIZE INPUT_CHANNEL*STEP1_CONTAINER_INPUT_H*STEP1_CONTAINER_INPUT_W //input
|
|
#define STEP1_CONTAINER_OUTPUT_SIZE INPUT_CHANNEL*(STEP1_CONTAINER_BBOX_SIZE1*STEP1_CONTAINER_BBOX_SIZE1+STEP1_CONTAINER_BBOX_SIZE2*STEP1_CONTAINER_BBOX_SIZE2+STEP1_CONTAINER_BBOX_SIZE3*STEP1_CONTAINER_BBOX_SIZE3)*(STEP1_CONTAINER_CLEAR_NUM+STEP1_CONTAINER_CLASS_NUM+STEP1_CONTAINER_POS_CONF) //output
|
|
|
|
// == step 1 👆 yolov5 ==
|
|
// == step 2 👇 yolov5 ==
|
|
// -- 车号/属性 --
|
|
#define STEP2_INPUT_BLOB_NAME "images" //输入层名称
|
|
#define STEP2_OUTPUT_BLOB_NAME "output" //输出层名称
|
|
|
|
#define STEP2_INPUT_CHANNEL 3 //输入通道数
|
|
#define STEP2_BATCH_SIZE 1 //batch size
|
|
|
|
#define STEP2_POS_CONF 5 //置信度加坐标
|
|
#define STEP2_LOCATIONS 4 //坐标数目
|
|
|
|
#define STEP2_NMS_THRESH 0.5 //step2 非极大值抑制阈值
|
|
#define STEP2_SCORE_THRESH 0.6 //step2 置信度(得分)阈值
|
|
|
|
#define STEP2_CLASS_NUM 47 //step2 分类数
|
|
#define STEP2_CLEAR_NUM 5 //step2 清晰度个数
|
|
|
|
#define STEP2_INPUT_H 608 //step2 输入图像高
|
|
#define STEP2_INPUT_W 608 //step2 输入图像宽
|
|
#define STEP2_BATCH_SIZE 1
|
|
#define STEP2_OUTPUT_HISTOGRAM_N 5 //step2 不能超过22743
|
|
|
|
#define STEP2_BBOX_SIZE1 (STEP2_INPUT_H/32)
|
|
#define STEP2_BBOX_SIZE2 (STEP2_INPUT_H/16)
|
|
#define STEP2_BBOX_SIZE3 (STEP2_INPUT_H/8)
|
|
|
|
#define STEP2_DET_SIZE (STEP2_CLASS_NUM + STEP2_CLEAR_NUM + STEP2_POS_CONF)
|
|
|
|
#define STEP2_INPUT_SIZE INPUT_CHANNEL*STEP2_INPUT_H*STEP2_INPUT_W //input
|
|
#define STEP2_OUTPUT_SIZE INPUT_CHANNEL*(STEP2_BBOX_SIZE1*STEP2_BBOX_SIZE1+STEP2_BBOX_SIZE2*STEP2_BBOX_SIZE2+STEP2_BBOX_SIZE3*STEP2_BBOX_SIZE3)*(STEP2_CLEAR_NUM+STEP2_CLASS_NUM+STEP2_POS_CONF) //output
|
|
// -- 集装箱 --
|
|
#define STEP2_CONTAINER_INPUT_BLOB_NAME "images" //输入层名称
|
|
#define STEP2_CONTAINER_OUTPUT_BLOB_NAME "output" //输出层名称
|
|
#define STEP2_CONTAINER_INPUT_CHANNEL 3 //输入通道数
|
|
#define STEP2_CONTAINER_BATCH_SIZE 1 //batch size
|
|
#define STEP2_CONTAINER_POS_CONF 5 //置信度加坐标
|
|
#define STEP2_CONTAINER_LOCATIONS 4 //坐标数目
|
|
#define STEP2_CONTAINER_NMS_THRESH 0.5 //step1 非极大值抑制阈值
|
|
#define STEP2_CONTAINER_SCORE_THRESH 0.9 //step1 置信度(得分)阈值
|
|
#define STEP2_CONTAINER_CLASS_NUM 7 //step1 分类数
|
|
#define STEP2_CONTAINER_CLEAR_NUM 21 //step1 清晰度个数
|
|
#define STEP2_CONTAINER_INPUT_H 960 //step1 输入图像高
|
|
#define STEP2_CONTAINER_INPUT_W 960 //step1 输入图像宽
|
|
#define STEP2_CONTAINER_BATCH_SIZE 1
|
|
#define STEP2_CONTAINER_BBOX_SIZE1 (STEP2_CONTAINER_INPUT_H/32)
|
|
#define STEP2_CONTAINER_BBOX_SIZE2 (STEP2_CONTAINER_INPUT_H/16)
|
|
#define STEP2_CONTAINER_BBOX_SIZE3 (STEP2_CONTAINER_INPUT_H/8)
|
|
#define STEP2_CONTAINER_DET_SIZE (STEP2_CONTAINER_CLASS_NUM + STEP2_CONTAINER_CLEAR_NUM + STEP2_CONTAINER_POS_CONF)
|
|
#define STEP2_CONTAINER_INPUT_SIZE INPUT_CHANNEL*STEP2_CONTAINER_INPUT_H*STEP2_CONTAINER_INPUT_W //input
|
|
#define STEP2_CONTAINER_OUTPUT_SIZE INPUT_CHANNEL*(STEP2_CONTAINER_BBOX_SIZE1*STEP2_CONTAINER_BBOX_SIZE1+STEP2_CONTAINER_BBOX_SIZE2*STEP2_CONTAINER_BBOX_SIZE2+STEP2_CONTAINER_BBOX_SIZE3*STEP2_CONTAINER_BBOX_SIZE3)*(STEP2_CONTAINER_CLEAR_NUM+STEP2_CONTAINER_CLASS_NUM+STEP2_CONTAINER_POS_CONF) //output
|
|
|
|
// == step 2 👆 yolov5 ==
|
|
|
|
|
|
// -- 来车检测 retinanet --
|
|
#define RETINANET_CLASSIFY_INPUT_CHANNEL 3 //输入通道数
|
|
#define RETINANET_CLASSIFY_BATCH_SIZE 1 //batch size
|
|
|
|
#define RETINANET_CLASSIFY_INPUT_H 537 // 模型输入高
|
|
#define RETINANET_CLASSIFY_INPUT_W 925 // 模型输入宽
|
|
|
|
#define RETINANET_CLASSIFY_INPUT_SIZE (RETINANET_CLASSIFY_BATCH_SIZE*RETINANET_CLASSIFY_INPUT_CHANNEL*RETINANET_CLASSIFY_INPUT_W*RETINANET_CLASSIFY_INPUT_H) //input
|
|
#define RETINANET_CLASSIFY_OUTPUT_SIZE 2 //output
|
|
|
|
#define RETINANET_CLASSIFY_INPUT_BLOB_NAME "data" //输入层名称
|
|
#define RETINANET_CLASSIFY_OUTPUT_BLOB_NAME "train" //输出层名称
|
|
|
|
|
|
typedef struct alignas(float) _Detection{
|
|
float fBbox[LOCATIONS];
|
|
float fClassConf;
|
|
int iClassId;
|
|
}Detection;
|
|
|
|
typedef struct alignas(float) _ClearDetection{
|
|
Detection detection;
|
|
float fClearConf;
|
|
int iClearId;
|
|
}ClearDetection;
|
|
|
|
typedef struct _ModelParam{
|
|
unsigned int uiClassNum;
|
|
|
|
unsigned int uiDetSize;
|
|
float fScoreThreshold;
|
|
float fNmsThreshold;
|
|
}ModelParam;
|
|
|
|
typedef struct _ClearModelParam{
|
|
unsigned int uiClearNum;
|
|
ModelParam modelParam;
|
|
}ClearModelParam;
|
|
|
|
typedef struct _ModelInfo{
|
|
unsigned int uiModelWidth;
|
|
unsigned int uiModelHeight;
|
|
unsigned int uiInputSize;
|
|
unsigned int uiOutputSize;
|
|
unsigned int uiChannel;
|
|
unsigned int uiBatchSize;
|
|
std::string strInputBlobName;
|
|
std::string strOutputBlobName;
|
|
}ModelInfo;
|
|
|
|
typedef struct _ClearModelInfo{
|
|
ClearModelParam clearModelParam;
|
|
ModelInfo modelInfo;
|
|
}ClearModelInfo;
|
|
|
|
typedef struct _CommonModelInfo{
|
|
ModelParam modelParam;
|
|
ModelInfo modelInfo;
|
|
}CommonModelInfo;
|
|
|
|
#endif //END OF _MODEL_H_
|