HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Update YOLOv8-ONNXRuntime-CPP example with GPU inference (#4328)

Browse source 
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>
This commit is contained in:
Onuralp SEZER 2023-08-13 18:34:39 +03:00 • committed by GitHub
parent f6b58e9d75
commit b5d1af42d8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 364 additions and 362 deletions
Download patch file Download diff file Expand all files Collapse all files

96
examples/YOLOv8-ONNXRuntime-CPP/inference.h
View file

@ -1,6 +1,6 @@
#pragma once
#define RET_OK nullptr
#define RET_OK nullptr
#ifdef _WIN32
#include <Windows.h>
@ -13,72 +13,72 @@
#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
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_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;
typedef struct _DCSP_RESULT {
int classId;
float confidence;
cv::Rect box;
} DCSP_RESULT;
class DCSP_CORE
{
class DCSP_CORE {
public:
DCSP_CORE();
~DCSP_CORE();
DCSP_CORE();
~DCSP_CORE();
public:
char* CreateSession(DCSP_INIT_PARAM &iParams);
char *CreateSession(DCSP_INIT_PARAM &iParams);
char* RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT>& oResult);
char *RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT> &oResult);
char* WarmUpSession();
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);
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;
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;
MODEL_TYPE modelType;
std::vector<int> imgSize;
float rectConfidenceThreshold;
float iouThreshold;
};