generated from zhangwei/Matrixai
1 line
3.3 KiB
C++
1 line
3.3 KiB
C++
#ifndef _YOLOV5_CLERAITY_INFERENCE_H_
|
|
#define _YOLOV5_CLERAITY_INFERENCE_H_
|
|
|
|
#include <string>
|
|
|
|
#include "logging.h"
|
|
#include <opencv2/opencv.hpp>
|
|
#include <NvCaffeParser.h>
|
|
#include <NvInfer.h>
|
|
#include <NvInferPlugin.h>
|
|
#include <NvOnnxParser.h>
|
|
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cuda_runtime_api.h>
|
|
|
|
typedef struct _YoloV5ModelParam
|
|
{
|
|
unsigned int uiClassNum;
|
|
unsigned int uiClearNum;
|
|
unsigned int uiDetSize;
|
|
float fScoreThreshold;
|
|
float fNmsThreshold;
|
|
} YoloV5ModelParam;
|
|
|
|
typedef struct _ModelCommonInfo
|
|
{
|
|
unsigned int uiModelWidth;
|
|
unsigned int uiModelHeight;
|
|
unsigned int uiInputSize;
|
|
unsigned int uiOutputSize;
|
|
unsigned int uiChannel;
|
|
unsigned int uiBatchSize;
|
|
std::string strInputBlobName;
|
|
std::string strOutputBlobName;
|
|
} ModelCommonInfo;
|
|
|
|
typedef struct _YoloV5ModelInfo
|
|
{
|
|
YoloV5ModelParam yolov5ModelParam;
|
|
ModelCommonInfo modelCommonInfo;
|
|
} YoloV5ModelInfo;
|
|
|
|
#define LOCATIONS 4
|
|
struct alignas(float) stDetection { //alignas(float)
|
|
float bbox[LOCATIONS];
|
|
float class_conf;
|
|
float clear_conf;
|
|
int class_id;
|
|
int clear_id;
|
|
};
|
|
|
|
class YoloV5Inference
|
|
{
|
|
public:
|
|
YoloV5Inference();
|
|
~YoloV5Inference();
|
|
|
|
int YoloV5InferenceInit(YoloV5ModelInfo &stYoloV5ModelInfo, std::string strModelName, std::string strEngineName);
|
|
int YoloV5InferenceDeinit(void);
|
|
int YoloV5InferenceModel(cv::Mat& frame, std::vector<stDetection>& res);
|
|
private:
|
|
float GetResizeRatio(unsigned int img_width, unsigned int img_height, unsigned int model_width, unsigned int model_height);
|
|
void doInference(nvinfer1::IExecutionContext& context, cudaStream_t& stream, void** buffers, unsigned int inputIndex, float* input, int inputSize,
|
|
unsigned int ouputIndex, float* output, int outputSize, int batchSize);
|
|
void doInferenceV2(nvinfer1::IExecutionContext& context, cudaStream_t& stream, void** buffers, unsigned int outputIndex, float* output, int outputSize, int batchSize);
|
|
void decode_opencv_nms(std::vector<stDetection>& res, float *output, int outSize, unsigned int detSize, unsigned int classNum, unsigned int clearNum, float confThresh, float nmsThresh = 0.5);
|
|
void CenterResetLocation(float fResizeRatio, unsigned int orig_width, unsigned int orig_height, unsigned int input_w, unsigned int input_h, stDetection &detection);
|
|
void xywh2xyxy(float *xywh, float * xyxy);
|
|
cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h);
|
|
private:
|
|
YoloV5ModelInfo stYoloV5ModelInfo_;
|
|
cudaStream_t* pImagePreprocessStream_ = nullptr; // 图像预处理CUDA流
|
|
cudaStream_t* pInferenceModelStream_ = nullptr; // 模型推理CUDA流
|
|
|
|
float* pfBuffers_[2];
|
|
float* pfInputData_ = nullptr;
|
|
float* pfOutputData_ = nullptr;
|
|
|
|
uint8_t* pu8ImgHost_ = nullptr; // 相关内存分配
|
|
uint8_t* pu8ImgDevice_ = nullptr;
|
|
|
|
unsigned int uiInputIndex_ = 0, uiOutputIndex_ = 0;
|
|
|
|
Logger* pGLogger_ = nullptr;
|
|
nvinfer1::IRuntime* pRuntime_ = nullptr;
|
|
nvinfer1::ICudaEngine* pEngine_ = nullptr;
|
|
nvinfer1::IExecutionContext* pContext_ = nullptr;
|
|
};
|
|
|
|
#endif // END OF _YOLOV5_CLERAITY_INFERENCE_H_
|