mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 13:34:23 +08:00
22474e9ad5
Signed-off-by: Onuralp SEZER <thunderbirdtr@fedoraproject.org> Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
81 lines
1.5 KiB
C++
81 lines
1.5 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"
|
|
|
|
|
|
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
|
|
};
|
|
|
|
|
|
|
|
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;
|
|
};
|