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.
68 lines
1.7 KiB
C++
68 lines
1.7 KiB
C++
#pragma once
|
|
|
|
#define RET_OK nullptr
|
|
|
|
#ifdef _WIN32
|
|
#include <Windows.h>
|
|
#include <direct.h>
|
|
#include <io.h>
|
|
#endif
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
#include <cstdio>
|
|
#include <opencv2/opencv.hpp>
|
|
#include <onnxruntime_cxx_api.h>
|
|
#include "YOLOManager.h"
|
|
|
|
#ifdef USE_CUDA
|
|
#include <cuda_fp16.h>
|
|
#endif
|
|
|
|
class YOLO_V8 :public YOLO
|
|
{
|
|
public:
|
|
YOLO_V8();
|
|
|
|
~YOLO_V8();
|
|
|
|
public:
|
|
bool InitParam(DL_INIT_PARAM& iParams)override;
|
|
|
|
bool CreateSession(std::string model_path, std::string model_name)override;
|
|
bool CreateSession(std::string model_path)override;
|
|
|
|
char* RunSession(cv::Mat& in, cv::Mat& out, std::vector<std::pair<int, cv::Rect> >& results)override;
|
|
char* RunSession(std::vector<cv::Mat>& vec_in, std::vector<cv::Mat>& vec_out, std::vector<std::vector<std::pair<int, cv::Rect> > >& vec_results)override;
|
|
|
|
char* WarmUpSession()override;
|
|
char* WarmUpSession(int shoot)override;
|
|
|
|
template<typename T>
|
|
char* TensorProcess(clock_t& starttime_1, cv::Mat& iImg, T* blob,
|
|
std::vector<int64_t>& inputNodeDims,
|
|
std::vector<DL_RESULT>& oResult);
|
|
|
|
char* PreProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg);
|
|
|
|
std::vector<std::string> classes{};
|
|
|
|
private:
|
|
Ort::Env env;
|
|
Ort::Session* session;
|
|
bool cudaEnable;
|
|
Ort::RunOptions options;
|
|
//std::vector<const char*> inputNodeNames;
|
|
//std::vector<const char*> outputNodeNames;
|
|
int logSeverityLevel;
|
|
int intraOpNumThreads;
|
|
MODEL_TYPE modelType;
|
|
std::vector<int> imgSize;
|
|
float rectConfidenceThreshold;
|
|
float iouThreshold;
|
|
float resizeScales;//letterbox scale
|
|
std::vector<char*> inputNodeNames;
|
|
std::vector<char*> outputNodeNames;
|
|
};
|
|
|