File size: 1,744 Bytes
ab854b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
#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"
#ifdef USE_CUDA
#include <cuda_fp16.h>
#endif
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;
};
|