You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

43 lines
1.1 KiB
C++

#pragma once
#include "MF_ObjectDetectBase.h"
class MF_Yolov8Infer : public MF_ObjectDetectBase
{
public:
MF_Yolov8Infer(const trtUtils::InitParameter& param);
~MF_Yolov8Infer();
// 初始化引擎 engine
bool initEngine(const std::string& _onnxFileName);
// 推理
bool doTRTInfer(const std::vector<MN_VisionImage::MS_ImageParam>& _bufImgs, std::vector<trtUtils::MR_Result>* _detectRes, int* _user);
// 推理
bool doTRTInfer(const std::vector<cv::Mat>& _matImgs, std::vector<trtUtils::MR_Result>* _detectRes, int* _user);
// 获取错误信息
static std::string getError();
// 清理数据/内存
void freeMemeory();
protected:
// 将图像数据拷贝到显存
int copyToDevice(const std::vector<cv::Mat>& _imgsBatch);
// 预处理
int preProcess(const std::vector<cv::Mat>& _imgsBatch);
// 推理
int infer();
// 将推理结果从显存拷贝到cpu
int copyFromDevice(const std::vector<cv::Mat>& _imgsBatch);
// 后处理
int postProcess(const std::vector<cv::Mat>& _imgsBatch);
// 获取最终检测结果
int getDetectResult(std::vector<trtUtils::MR_Result>& _result);
private:
float* m_output_src_transpose_device;
};