You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
84 lines
1.6 KiB
84 lines
1.6 KiB
#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 <cuda_fp16.h> |
|
|
|
|
|
enum MODEL_TYPE |
|
{ |
|
//FLOAT32 MODEL |
|
YOLO_ORIGIN_V5 = 0, |
|
YOLO_ORIGIN_V8 = 1,//only support v8 detector currently |
|
YOLO_POSE_V8 = 2, |
|
YOLO_CLS_V8 = 3, |
|
YOLO_ORIGIN_V8_HALF = 4, |
|
YOLO_POSE_V8_HALF = 5, |
|
YOLO_CLS_V8_HALF = 6 |
|
}; |
|
|
|
|
|
|
|
typedef struct _DCSP_INIT_PARAM |
|
{ |
|
std::string ModelPath; |
|
MODEL_TYPE ModelType = YOLO_ORIGIN_V8; |
|
std::vector<int> imgSize={640, 640}; |
|
float RectConfidenceThreshold = 0.6; |
|
float iouThreshold = 0.5; |
|
bool CudaEnable = false; |
|
int LogSeverityLevel = 3; |
|
int IntraOpNumThreads = 1; |
|
}DCSP_INIT_PARAM; |
|
|
|
|
|
typedef struct _DCSP_RESULT |
|
{ |
|
int classId; |
|
float confidence; |
|
cv::Rect box; |
|
}DCSP_RESULT; |
|
|
|
|
|
class DCSP_CORE |
|
{ |
|
public: |
|
DCSP_CORE(); |
|
~DCSP_CORE(); |
|
|
|
public: |
|
char* CreateSession(DCSP_INIT_PARAM &iParams); |
|
|
|
char* RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT>& oResult); |
|
|
|
char* WarmUpSession(); |
|
|
|
template<typename N> |
|
char* TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims, std::vector<DCSP_RESULT>& oResult); |
|
|
|
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; |
|
|
|
MODEL_TYPE modelType; |
|
std::vector<int> imgSize; |
|
float rectConfidenceThreshold; |
|
float iouThreshold; |
|
};
|
|
|