Update YOLOv8-ONNXRuntime-CPP example with GPU inference (#4328)
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
GitHub
@ -5,14 +5,14 @@ This repository features a collection of real-world applications and walkthrough
|
||||
### Ultralytics YOLO Example Applications
|
||||
|
||||
| Title | Format | Contributor |
|
||||
| -------------------------------------------------------------------------------------------------------------- | ------------------ | --------------------------------------------------- |
|
||||
| -------------------------------------------------------------------------------------------------------------- | ------------------ | ----------------------------------------------------------------------------------------- |
|
||||
| [YOLO ONNX Detection Inference with C++](./YOLOv8-CPP-Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) |
|
||||
| [YOLO OpenCV ONNX Detection Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) |
|
||||
| [YOLOv8 .NET ONNX ImageSharp](https://github.com/dme-compunet/YOLOv8) | C#/ONNX/ImageSharp | [Compunet](https://github.com/dme-compunet) |
|
||||
| [YOLO .Net ONNX Detection C#](https://www.nuget.org/packages/Yolov8.Net) | C# .Net | [Samuel Stainback](https://github.com/sstainba) |
|
||||
| [YOLOv8 on NVIDIA Jetson(TensorRT and DeepStream)](https://wiki.seeedstudio.com/YOLOv8-DeepStream-TRT-Jetson/) | Python | [Lakshantha](https://github.com/lakshanthad) |
|
||||
| [YOLOv8 ONNXRuntime Python](./YOLOv8-ONNXRuntime) | Python/ONNXRuntime | [Semih Demirel](https://github.com/semihhdemirel) |
|
||||
| [YOLOv8-ONNXRuntime-CPP](./YOLOv8-ONNXRuntime-CPP) | C++/ONNXRuntime | [DennisJcy](https://github.com/DennisJcy) |
|
||||
| [YOLOv8-ONNXRuntime-CPP](./YOLOv8-ONNXRuntime-CPP) | C++/ONNXRuntime | [DennisJcy](https://github.com/DennisJcy), [Onuralp Sezer](https://github.com/onuralpszr) |
|
||||
| [RTDETR ONNXRuntime C#](https://github.com/Kayzwer/yolo-cs/blob/master/RTDETR.cs) | C#/ONNX | [Kayzwer](https://github.com/Kayzwer) |
|
||||
|
||||
### How to Contribute
|
||||
|
||||
@ -17,9 +17,12 @@ include_directories(${OpenCV_INCLUDE_DIRS})
|
||||
|
||||
|
||||
# -------------- Compile CUDA for FP16 inference if needed ------------------#
|
||||
option(USE_CUDA "Enable CUDA support" ON)
|
||||
if (USE_CUDA)
|
||||
find_package(CUDA REQUIRED)
|
||||
include_directories(${CUDA_INCLUDE_DIRS})
|
||||
|
||||
add_definitions(-DUSE_CUDA)
|
||||
endif ()
|
||||
|
||||
# ONNXRUNTIME
|
||||
|
||||
@ -27,15 +30,17 @@ include_directories(${CUDA_INCLUDE_DIRS})
|
||||
set(ONNXRUNTIME_VERSION 1.15.1)
|
||||
|
||||
if (WIN32)
|
||||
# CPU
|
||||
# set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-win-x64-${ONNXRUNTIME_VERSION}")
|
||||
# GPU
|
||||
if (USE_CUDA)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-win-x64-gpu-${ONNXRUNTIME_VERSION}")
|
||||
else ()
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-win-x64-${ONNXRUNTIME_VERSION}")
|
||||
endif ()
|
||||
elseif (LINUX)
|
||||
# CPU
|
||||
# set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}")
|
||||
# GPU
|
||||
if (USE_CUDA)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-linux-x64-gpu-${ONNXRUNTIME_VERSION}")
|
||||
else ()
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}")
|
||||
endif ()
|
||||
elseif (APPLE)
|
||||
set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/onnxruntime-osx-arm64-${ONNXRUNTIME_VERSION}")
|
||||
# Apple X64 binary
|
||||
@ -55,9 +60,15 @@ set(PROJECT_SOURCES
|
||||
add_executable(${PROJECT_NAME} ${PROJECT_SOURCES})
|
||||
|
||||
if (WIN32)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/onnxruntime.lib ${CUDA_LIBRARIES})
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/onnxruntime.lib)
|
||||
if (USE_CUDA)
|
||||
target_link_libraries(${PROJECT_NAME} ${CUDA_LIBRARIES})
|
||||
endif ()
|
||||
elseif (LINUX)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.so ${CUDA_LIBRARIES})
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.so)
|
||||
if (USE_CUDA)
|
||||
target_link_libraries(${PROJECT_NAME} ${CUDA_LIBRARIES})
|
||||
endif ()
|
||||
elseif (APPLE)
|
||||
target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.dylib)
|
||||
endif ()
|
||||
|
||||
@ -28,16 +28,23 @@ Alternatively, you can use the following command for exporting the model in the
|
||||
yolo export model=yolov8n.pt opset=12 simplify=True dynamic=False format=onnx imgsz=640,640
|
||||
```
|
||||
|
||||
## Download COCO.yaml file
|
||||
|
||||
In order to run example, you also need to download coco.yaml. You can download the file manually from [here](https://raw.githubusercontent.com/ultralytics/ultralytics/main/ultralytics/cfg/datasets/coco.yaml)
|
||||
|
||||
## Dependencies
|
||||
|
||||
| Dependency | Version |
|
||||
| -------------------------------- | -------- |
|
||||
| -------------------------------- | ------------- |
|
||||
| Onnxruntime(linux,windows,macos) | >=1.14.1 |
|
||||
| OpenCV | >=4.0.0 |
|
||||
| C++ | >=17 |
|
||||
| Cmake | >=3.5 |
|
||||
| Cuda (Optional) | >=11.4,\<12.0 |
|
||||
| cuDNN (Cuda required) | =8 |
|
||||
|
||||
Note: The dependency on C++17 is due to the usage of the C++17 filesystem feature.
|
||||
Note (2): Due to ONNX Runtime, we need to use CUDA 11 and cuDNN 8. Keep in mind that this requirement might change in the future.
|
||||
|
||||
## Usage
|
||||
|
||||
|
||||
@ -3,39 +3,35 @@
|
||||
|
||||
#define benchmark
|
||||
|
||||
DCSP_CORE::DCSP_CORE()
|
||||
{
|
||||
DCSP_CORE::DCSP_CORE() {
|
||||
|
||||
}
|
||||
|
||||
|
||||
DCSP_CORE::~DCSP_CORE()
|
||||
{
|
||||
DCSP_CORE::~DCSP_CORE() {
|
||||
delete session;
|
||||
}
|
||||
|
||||
|
||||
#ifdef USE_CUDA
|
||||
namespace Ort
|
||||
{
|
||||
template<>
|
||||
struct TypeToTensorType<half> { static constexpr ONNXTensorElementDataType type = ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; };
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
template<typename T>
|
||||
char* BlobFromImage(cv::Mat& iImg, T& iBlob)
|
||||
{
|
||||
char *BlobFromImage(cv::Mat &iImg, T &iBlob) {
|
||||
int channels = iImg.channels();
|
||||
int imgHeight = iImg.rows;
|
||||
int imgWidth = iImg.cols;
|
||||
|
||||
for (int c = 0; c < channels; c++)
|
||||
{
|
||||
for (int h = 0; h < imgHeight; h++)
|
||||
{
|
||||
for (int w = 0; w < imgWidth; w++)
|
||||
{
|
||||
iBlob[c * imgWidth * imgHeight + h * imgWidth + w] = typename std::remove_pointer<T>::type((iImg.at<cv::Vec3b>(h, w)[c]) / 255.0f);
|
||||
for (int c = 0; c < channels; c++) {
|
||||
for (int h = 0; h < imgHeight; h++) {
|
||||
for (int w = 0; w < imgWidth; w++) {
|
||||
iBlob[c * imgWidth * imgHeight + h * imgWidth + w] = typename std::remove_pointer<T>::type(
|
||||
(iImg.at<cv::Vec3b>(h, w)[c]) / 255.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -43,12 +39,10 @@ char* BlobFromImage(cv::Mat& iImg, T& iBlob)
|
||||
}
|
||||
|
||||
|
||||
char* PostProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg)
|
||||
{
|
||||
char *PostProcess(cv::Mat &iImg, std::vector<int> iImgSize, cv::Mat &oImg) {
|
||||
cv::Mat img = iImg.clone();
|
||||
cv::resize(iImg, oImg, cv::Size(iImgSize.at(0), iImgSize.at(1)));
|
||||
if (img.channels() == 1)
|
||||
{
|
||||
if (img.channels() == 1) {
|
||||
cv::cvtColor(oImg, oImg, cv::COLOR_GRAY2BGR);
|
||||
}
|
||||
cv::cvtColor(oImg, oImg, cv::COLOR_BGR2RGB);
|
||||
@ -56,27 +50,23 @@ char* PostProcess(cv::Mat& iImg, std::vector<int> iImgSize, cv::Mat& oImg)
|
||||
}
|
||||
|
||||
|
||||
char* DCSP_CORE::CreateSession(DCSP_INIT_PARAM &iParams)
|
||||
{
|
||||
char *DCSP_CORE::CreateSession(DCSP_INIT_PARAM &iParams) {
|
||||
char *Ret = RET_OK;
|
||||
std::regex pattern("[\u4e00-\u9fa5]");
|
||||
bool result = std::regex_search(iParams.ModelPath, pattern);
|
||||
if (result)
|
||||
{
|
||||
if (result) {
|
||||
Ret = "[DCSP_ONNX]:Model path error.Change your model path without chinese characters.";
|
||||
std::cout << Ret << std::endl;
|
||||
return Ret;
|
||||
}
|
||||
try
|
||||
{
|
||||
try {
|
||||
rectConfidenceThreshold = iParams.RectConfidenceThreshold;
|
||||
iouThreshold = iParams.iouThreshold;
|
||||
imgSize = iParams.imgSize;
|
||||
modelType = iParams.ModelType;
|
||||
env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "Yolo");
|
||||
Ort::SessionOptions sessionOption;
|
||||
if (iParams.CudaEnable)
|
||||
{
|
||||
if (iParams.CudaEnable) {
|
||||
cudaEnable = iParams.CudaEnable;
|
||||
OrtCUDAProviderOptions cudaOption;
|
||||
cudaOption.device_id = 0;
|
||||
@ -99,16 +89,14 @@ char* DCSP_CORE::CreateSession(DCSP_INIT_PARAM &iParams)
|
||||
session = new Ort::Session(env, modelPath, sessionOption);
|
||||
Ort::AllocatorWithDefaultOptions allocator;
|
||||
size_t inputNodesNum = session->GetInputCount();
|
||||
for (size_t i = 0; i < inputNodesNum; i++)
|
||||
{
|
||||
for (size_t i = 0; i < inputNodesNum; i++) {
|
||||
Ort::AllocatedStringPtr input_node_name = session->GetInputNameAllocated(i, allocator);
|
||||
char *temp_buf = new char[50];
|
||||
strcpy(temp_buf, input_node_name.get());
|
||||
inputNodeNames.push_back(temp_buf);
|
||||
}
|
||||
size_t OutputNodesNum = session->GetOutputCount();
|
||||
for (size_t i = 0; i < OutputNodesNum; i++)
|
||||
{
|
||||
for (size_t i = 0; i < OutputNodesNum; i++) {
|
||||
Ort::AllocatedStringPtr output_node_name = session->GetOutputNameAllocated(i, allocator);
|
||||
char *temp_buf = new char[10];
|
||||
strcpy(temp_buf, output_node_name.get());
|
||||
@ -118,8 +106,7 @@ char* DCSP_CORE::CreateSession(DCSP_INIT_PARAM &iParams)
|
||||
WarmUpSession();
|
||||
return RET_OK;
|
||||
}
|
||||
catch (const std::exception& e)
|
||||
{
|
||||
catch (const std::exception &e) {
|
||||
const char *str1 = "[DCSP_ONNX]:";
|
||||
const char *str2 = e.what();
|
||||
std::string result = std::string(str1) + std::string(str2);
|
||||
@ -133,8 +120,7 @@ char* DCSP_CORE::CreateSession(DCSP_INIT_PARAM &iParams)
|
||||
}
|
||||
|
||||
|
||||
char* DCSP_CORE::RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT>& oResult)
|
||||
{
|
||||
char *DCSP_CORE::RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT> &oResult) {
|
||||
#ifdef benchmark
|
||||
clock_t starttime_1 = clock();
|
||||
#endif // benchmark
|
||||
@ -142,19 +128,18 @@ char* DCSP_CORE::RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT>& oResult)
|
||||
char *Ret = RET_OK;
|
||||
cv::Mat processedImg;
|
||||
PostProcess(iImg, imgSize, processedImg);
|
||||
if (modelType < 4)
|
||||
{
|
||||
if (modelType < 4) {
|
||||
float *blob = new float[processedImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> inputNodeDims = {1, 3, imgSize.at(0), imgSize.at(1)};
|
||||
TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult);
|
||||
}
|
||||
else
|
||||
{
|
||||
} else {
|
||||
#ifdef USE_CUDA
|
||||
half* blob = new half[processedImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> inputNodeDims = { 1,3,imgSize.at(0),imgSize.at(1) };
|
||||
TensorProcess(starttime_1, iImg, blob, inputNodeDims, oResult);
|
||||
#endif
|
||||
}
|
||||
|
||||
return Ret;
|
||||
@ -162,13 +147,16 @@ char* DCSP_CORE::RunSession(cv::Mat &iImg, std::vector<DCSP_RESULT>& oResult)
|
||||
|
||||
|
||||
template<typename N>
|
||||
char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std::vector<int64_t>& inputNodeDims, std::vector<DCSP_RESULT>& oResult)
|
||||
{
|
||||
Ort::Value inputTensor = Ort::Value::CreateTensor<typename std::remove_pointer<N>::type>(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), inputNodeDims.data(), inputNodeDims.size());
|
||||
char *DCSP_CORE::TensorProcess(clock_t &starttime_1, cv::Mat &iImg, N &blob, std::vector<int64_t> &inputNodeDims,
|
||||
std::vector<DCSP_RESULT> &oResult) {
|
||||
Ort::Value inputTensor = Ort::Value::CreateTensor<typename std::remove_pointer<N>::type>(
|
||||
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1),
|
||||
inputNodeDims.data(), inputNodeDims.size());
|
||||
#ifdef benchmark
|
||||
clock_t starttime_2 = clock();
|
||||
#endif // benchmark
|
||||
auto outputTensor = session->Run(options, inputNodeNames.data(), &inputTensor, 1, outputNodeNames.data(), outputNodeNames.size());
|
||||
auto outputTensor = session->Run(options, inputNodeNames.data(), &inputTensor, 1, outputNodeNames.data(),
|
||||
outputNodeNames.size());
|
||||
#ifdef benchmark
|
||||
clock_t starttime_3 = clock();
|
||||
#endif // benchmark
|
||||
@ -178,8 +166,7 @@ char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std
|
||||
std::vector<int64_t> outputNodeDims = tensor_info.GetShape();
|
||||
auto output = outputTensor.front().GetTensorMutableData<typename std::remove_pointer<N>::type>();
|
||||
delete blob;
|
||||
switch (modelType)
|
||||
{
|
||||
switch (modelType) {
|
||||
case 1://V8_ORIGIN_FP32
|
||||
case 4://V8_ORIGIN_FP16
|
||||
{
|
||||
@ -195,15 +182,13 @@ char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std
|
||||
|
||||
float x_factor = iImg.cols / 640.;
|
||||
float y_factor = iImg.rows / 640.;
|
||||
for (int i = 0; i < strideNum; ++i)
|
||||
{
|
||||
for (int i = 0; i < strideNum; ++i) {
|
||||
float *classesScores = data + 4;
|
||||
cv::Mat scores(1, this->classes.size(), CV_32FC1, classesScores);
|
||||
cv::Point class_id;
|
||||
double maxClassScore;
|
||||
cv::minMaxLoc(scores, 0, &maxClassScore, 0, &class_id);
|
||||
if (maxClassScore > rectConfidenceThreshold)
|
||||
{
|
||||
if (maxClassScore > rectConfidenceThreshold) {
|
||||
confidences.push_back(maxClassScore);
|
||||
class_ids.push_back(class_id.x);
|
||||
|
||||
@ -226,8 +211,7 @@ char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std
|
||||
std::vector<int> nmsResult;
|
||||
cv::dnn::NMSBoxes(boxes, confidences, rectConfidenceThreshold, iouThreshold, nmsResult);
|
||||
|
||||
for (int i = 0; i < nmsResult.size(); ++i)
|
||||
{
|
||||
for (int i = 0; i < nmsResult.size(); ++i) {
|
||||
int idx = nmsResult[i];
|
||||
DCSP_RESULT result;
|
||||
result.classId = class_ids[idx];
|
||||
@ -242,13 +226,12 @@ char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std
|
||||
double pre_process_time = (double) (starttime_2 - starttime_1) / CLOCKS_PER_SEC * 1000;
|
||||
double process_time = (double) (starttime_3 - starttime_2) / CLOCKS_PER_SEC * 1000;
|
||||
double post_process_time = (double) (starttime_4 - starttime_3) / CLOCKS_PER_SEC * 1000;
|
||||
if (cudaEnable)
|
||||
{
|
||||
std::cout << "[DCSP_ONNX(CUDA)]: " << pre_process_time << "ms pre-process, " << process_time << "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "[DCSP_ONNX(CPU)]: " << pre_process_time << "ms pre-process, " << process_time << "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
if (cudaEnable) {
|
||||
std::cout << "[DCSP_ONNX(CUDA)]: " << pre_process_time << "ms pre-process, " << process_time
|
||||
<< "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
} else {
|
||||
std::cout << "[DCSP_ONNX(CPU)]: " << pre_process_time << "ms pre-process, " << process_time
|
||||
<< "ms inference, " << post_process_time << "ms post-process." << std::endl;
|
||||
}
|
||||
#endif // benchmark
|
||||
|
||||
@ -259,29 +242,28 @@ char* DCSP_CORE::TensorProcess(clock_t& starttime_1, cv::Mat& iImg, N& blob, std
|
||||
}
|
||||
|
||||
|
||||
char* DCSP_CORE::WarmUpSession()
|
||||
{
|
||||
char *DCSP_CORE::WarmUpSession() {
|
||||
clock_t starttime_1 = clock();
|
||||
cv::Mat iImg = cv::Mat(cv::Size(imgSize.at(0), imgSize.at(1)), CV_8UC3);
|
||||
cv::Mat processedImg;
|
||||
PostProcess(iImg, imgSize, processedImg);
|
||||
if (modelType < 4)
|
||||
{
|
||||
if (modelType < 4) {
|
||||
float *blob = new float[iImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> YOLO_input_node_dims = {1, 3, imgSize.at(0), imgSize.at(1)};
|
||||
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1), YOLO_input_node_dims.data(), YOLO_input_node_dims.size());
|
||||
auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(), outputNodeNames.size());
|
||||
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
|
||||
Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU), blob, 3 * imgSize.at(0) * imgSize.at(1),
|
||||
YOLO_input_node_dims.data(), YOLO_input_node_dims.size());
|
||||
auto output_tensors = session->Run(options, inputNodeNames.data(), &input_tensor, 1, outputNodeNames.data(),
|
||||
outputNodeNames.size());
|
||||
delete[] blob;
|
||||
clock_t starttime_4 = clock();
|
||||
double post_process_time = (double) (starttime_4 - starttime_1) / CLOCKS_PER_SEC * 1000;
|
||||
if (cudaEnable)
|
||||
{
|
||||
if (cudaEnable) {
|
||||
std::cout << "[DCSP_ONNX(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
} else {
|
||||
#ifdef USE_CUDA
|
||||
half* blob = new half[iImg.total() * 3];
|
||||
BlobFromImage(processedImg, blob);
|
||||
std::vector<int64_t> YOLO_input_node_dims = { 1,3,imgSize.at(0),imgSize.at(1) };
|
||||
@ -294,6 +276,7 @@ char* DCSP_CORE::WarmUpSession()
|
||||
{
|
||||
std::cout << "[DCSP_ONNX(CUDA)]: " << "Cuda warm-up cost " << post_process_time << " ms. " << std::endl;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
return RET_OK;
|
||||
}
|
||||
|
||||
@ -13,11 +13,13 @@
|
||||
#include <cstdio>
|
||||
#include <opencv2/opencv.hpp>
|
||||
#include "onnxruntime_cxx_api.h"
|
||||
|
||||
#ifdef USE_CUDA
|
||||
#include <cuda_fp16.h>
|
||||
#endif
|
||||
|
||||
|
||||
enum MODEL_TYPE
|
||||
{
|
||||
enum MODEL_TYPE {
|
||||
//FLOAT32 MODEL
|
||||
YOLO_ORIGIN_V5 = 0,
|
||||
YOLO_ORIGIN_V8 = 1,//only support v8 detector currently
|
||||
@ -29,9 +31,7 @@ enum MODEL_TYPE
|
||||
};
|
||||
|
||||
|
||||
|
||||
typedef struct _DCSP_INIT_PARAM
|
||||
{
|
||||
typedef struct _DCSP_INIT_PARAM {
|
||||
std::string ModelPath;
|
||||
MODEL_TYPE ModelType = YOLO_ORIGIN_V8;
|
||||
std::vector<int> imgSize = {640, 640};
|
||||
@ -43,18 +43,17 @@ typedef struct _DCSP_INIT_PARAM
|
||||
} DCSP_INIT_PARAM;
|
||||
|
||||
|
||||
typedef struct _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();
|
||||
|
||||
public:
|
||||
@ -65,7 +64,8 @@ public:
|
||||
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);
|
||||
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{};
|
||||
|
||||
|
||||
@ -3,42 +3,41 @@
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
|
||||
void file_iterator(DCSP_CORE*& p)
|
||||
{
|
||||
void file_iterator(DCSP_CORE *&p) {
|
||||
std::filesystem::path current_path = std::filesystem::current_path();
|
||||
std::filesystem::path imgs_path = current_path / "images";
|
||||
for (auto& i : std::filesystem::directory_iterator(imgs_path))
|
||||
{
|
||||
if (i.path().extension() == ".jpg" || i.path().extension() == ".png")
|
||||
{
|
||||
for (auto &i: std::filesystem::directory_iterator(imgs_path)) {
|
||||
if (i.path().extension() == ".jpg" || i.path().extension() == ".png" || i.path().extension() == ".jpeg") {
|
||||
std::string img_path = i.path().string();
|
||||
cv::Mat img = cv::imread(img_path);
|
||||
std::vector<DCSP_RESULT> res;
|
||||
p->RunSession(img, res);
|
||||
|
||||
for (auto & re : res)
|
||||
{
|
||||
cv::rectangle(img, re.box, cv::Scalar(0, 0 , 255), 3);
|
||||
std::string label = p->classes[re.classId];
|
||||
for (auto &re: res) {
|
||||
cv::RNG rng(cv::getTickCount());
|
||||
cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256));
|
||||
|
||||
cv::rectangle(img, re.box, color, 3);
|
||||
std::string label = p->classes[re.classId] + " " + std::to_string(re.confidence);
|
||||
cv::putText(
|
||||
img,
|
||||
label,
|
||||
cv::Point(re.box.x, re.box.y - 5),
|
||||
cv::FONT_HERSHEY_SIMPLEX,
|
||||
0.75,
|
||||
cv::Scalar(255, 255, 0),
|
||||
color,
|
||||
2
|
||||
);
|
||||
}
|
||||
cv::imshow("Result", img);
|
||||
std::cout << "Press any key to exit" << std::endl;
|
||||
cv::imshow("Result of Detection", img);
|
||||
cv::waitKey(0);
|
||||
cv::destroyAllWindows();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int read_coco_yaml(DCSP_CORE*& p)
|
||||
{
|
||||
int read_coco_yaml(DCSP_CORE *&p) {
|
||||
// Open the YAML file
|
||||
std::ifstream file("coco.yaml");
|
||||
if (!file.is_open()) {
|
||||
@ -80,17 +79,19 @@ int read_coco_yaml(DCSP_CORE*& p)
|
||||
}
|
||||
|
||||
|
||||
int main()
|
||||
{
|
||||
int main() {
|
||||
DCSP_CORE *yoloDetector = new DCSP_CORE;
|
||||
std::string model_path = "yolov8n.onnx";
|
||||
read_coco_yaml(yoloDetector);
|
||||
#ifdef USE_CUDA
|
||||
// GPU FP32 inference
|
||||
DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8, {640, 640}, 0.1, 0.5, true };
|
||||
// GPU FP16 inference
|
||||
// DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8_HALF, {640, 640}, 0.1, 0.5, true };
|
||||
#else
|
||||
// CPU inference
|
||||
// DCSP_INIT_PARAM params{ model_path, YOLO_ORIGIN_V8, {640, 640}, 0.1, 0.5, false };
|
||||
DCSP_INIT_PARAM params{model_path, YOLO_ORIGIN_V8, {640, 640}, 0.1, 0.5, false};
|
||||
#endif
|
||||
yoloDetector->CreateSession(params);
|
||||
file_iterator(yoloDetector);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user