217 lines
8.7 KiB
C++
217 lines
8.7 KiB
C++
// YoloV8_ONNX.cpp (最终优化版)
|
||
|
||
#include "pch.h"
|
||
#include "Yolo_ONNX.h"
|
||
|
||
#include <onnxruntime_c_api.h>
|
||
#include <onnxruntime_cxx_api.h>
|
||
#include <opencv2/opencv.hpp>
|
||
#include <vector>
|
||
#include <string>
|
||
#include <iostream>
|
||
#include <memory>
|
||
|
||
namespace {
|
||
|
||
// 【优化】预处理函数
|
||
cv::Mat preprocess(const cv::Mat& img, int target_width, int target_height, int& pad_w, int& pad_h, float& scale) {
|
||
cv::Mat resized_img;
|
||
int w = img.cols;
|
||
int h = img.rows;
|
||
scale = std::min(static_cast<float>(target_width) / w, static_cast<float>(target_height) / h);
|
||
int new_w = static_cast<int>(w * scale);
|
||
int new_h = static_cast<int>(h * scale);
|
||
|
||
// 【优化】使用 INTER_AREA 插值算法，更适合图像缩小，与主流Python库行为更接近
|
||
cv::resize(img, resized_img, cv::Size(new_w, new_h), 0, 0, cv::INTER_AREA);
|
||
|
||
pad_w = target_width - new_w;
|
||
pad_h = target_height - new_h;
|
||
cv::Mat padded_img;
|
||
cv::copyMakeBorder(resized_img, padded_img, 0, pad_h, 0, pad_w, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
|
||
return padded_img;
|
||
}
|
||
|
||
// 后处理函数
|
||
std::vector<Detection> postprocess(Ort::Value& output_tensor, float scale, int pad_w, int pad_h, int img_w, int img_h, float conf_threshold, float iou_threshold) {
|
||
const auto output_shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();
|
||
const float* raw_output = output_tensor.GetTensorData<float>();
|
||
int num_classes = static_cast<int>(output_shape[1]) - 4;
|
||
int num_proposals = static_cast<int>(output_shape[2]);
|
||
|
||
std::vector<cv::Rect> boxes;
|
||
std::vector<float> scores;
|
||
std::vector<int> class_ids;
|
||
|
||
cv::Mat raw_data_mat(num_classes + 4, num_proposals, CV_32F, (void*)raw_output);
|
||
raw_data_mat = raw_data_mat.t();
|
||
|
||
for (int i = 0; i < num_proposals; ++i) {
|
||
const float* proposal = raw_data_mat.ptr<float>(i);
|
||
const float* class_scores = proposal + 4;
|
||
float max_score = 0.0f;
|
||
int class_id = -1;
|
||
for (int j = 0; j < num_classes; ++j) {
|
||
if (class_scores[j] > max_score) {
|
||
max_score = class_scores[j];
|
||
class_id = j;
|
||
}
|
||
}
|
||
if (max_score > conf_threshold) {
|
||
float cx = proposal[0];
|
||
float cy = proposal[1];
|
||
float w = proposal[2];
|
||
float h = proposal[3];
|
||
int left = static_cast<int>((cx - w / 2 - (pad_w / 2.0f)) / scale);
|
||
int top = static_cast<int>((cy - h / 2 - (pad_h / 2.0f)) / scale);
|
||
int width = static_cast<int>(w / scale);
|
||
int height = static_cast<int>(h / scale);
|
||
left = std::max(0, std::min(left, img_w - 1));
|
||
top = std::max(0, std::min(top, img_h - 1));
|
||
width = std::min(width, img_w - left);
|
||
height = std::min(height, img_h - top);
|
||
boxes.push_back(cv::Rect(left, top, width, height));
|
||
scores.push_back(max_score);
|
||
class_ids.push_back(class_id);
|
||
}
|
||
}
|
||
|
||
std::vector<int> nms_result;
|
||
cv::dnn::NMSBoxes(boxes, scores, conf_threshold, iou_threshold, nms_result);
|
||
std::vector<Detection> detections;
|
||
for (int idx : nms_result) {
|
||
detections.push_back({ class_ids[idx], scores[idx], boxes[idx].x, boxes[idx].y, boxes[idx].width, boxes[idx].height });
|
||
}
|
||
return detections;
|
||
}
|
||
}
|
||
|
||
|
||
extern "C" {
|
||
// 【修改】函数签名更新，增加了四个新参数
|
||
YOLO_API int perform_detection(
|
||
const wchar_t* model_path,
|
||
unsigned char* image_bytes,
|
||
int image_width,
|
||
int image_height,
|
||
Detection** out_detections,
|
||
int* out_detections_count,
|
||
const char** class_names,
|
||
int class_names_count,
|
||
float conf_threshold, // 使用传入的置信度阈值
|
||
float iou_threshold, // 使用传入的IOU阈值
|
||
int input_width, // 使用传入的模型输入宽度
|
||
int input_height // 使用传入的模型输入高度
|
||
) {
|
||
static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "YOLOv8-ONNX-GPU");
|
||
static std::unique_ptr<Ort::Session> session = nullptr;
|
||
// 【新增】用于判断模型是否需要重新加载的变量
|
||
static std::wstring current_model_path = L"";
|
||
|
||
try {
|
||
// 如果模型路径发生变化，则重新创建Session
|
||
if (!session || current_model_path != model_path) {
|
||
Ort::SessionOptions session_options;
|
||
OrtCUDAProviderOptions cuda_options;
|
||
session_options.AppendExecutionProvider_CUDA(cuda_options);
|
||
session = std::make_unique<Ort::Session>(env, model_path, session_options);
|
||
current_model_path = model_path;
|
||
}
|
||
|
||
// 【修改】移除硬编码的尺寸，使用接口传入的参数
|
||
std::vector<int64_t> input_shape = { 1, 3, input_height, input_width };
|
||
|
||
Ort::AllocatorWithDefaultOptions allocator;
|
||
std::string input_name_str = session->GetInputNameAllocated(0, allocator).get();
|
||
std::vector<const char*> input_node_names = { input_name_str.c_str() };
|
||
|
||
std::string output_name_str = session->GetOutputNameAllocated(0, allocator).get();
|
||
std::vector<const char*> output_node_names = { output_name_str.c_str() };
|
||
|
||
cv::Mat image(image_height, image_width, CV_8UC3, image_bytes);
|
||
if (image.empty()) return -1;
|
||
|
||
int pad_w, pad_h;
|
||
float scale;
|
||
// 【修改】使用接口传入的参数进行预处理
|
||
cv::Mat preprocessed_img = preprocess(image, input_width, input_height, pad_w, pad_h, scale);
|
||
|
||
cv::Mat blob;
|
||
cv::dnn::blobFromImage(preprocessed_img, blob, 1 / 255.0, cv::Size(), cv::Scalar(), false, false);
|
||
|
||
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||
Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_shape.data(), input_shape.size());
|
||
auto output_tensors = session->Run(Ort::RunOptions{ nullptr }, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1);
|
||
|
||
// 【修改】使用接口传入的参数进行后处理
|
||
std::vector<Detection> detections = postprocess(output_tensors[0], scale, pad_w, pad_h, image_width, image_height, conf_threshold, iou_threshold);
|
||
|
||
*out_detections_count = static_cast<int>(detections.size());
|
||
if (*out_detections_count > 0) {
|
||
*out_detections = new Detection[*out_detections_count];
|
||
std::copy(detections.begin(), detections.end(), *out_detections);
|
||
}
|
||
else {
|
||
*out_detections = nullptr;
|
||
}
|
||
}
|
||
catch (const Ort::Exception& e) {
|
||
std::cerr << "ONNX Runtime 异常: " << e.what() << std::endl;
|
||
return -2;
|
||
}
|
||
catch (const cv::Exception& e) {
|
||
std::cerr << "OpenCV 异常: " << e.what() << std::endl;
|
||
return -3;
|
||
}
|
||
catch (const std::exception& e) {
|
||
std::cerr << "标准库异常: " << e.what() << std::endl;
|
||
return -4;
|
||
}
|
||
return 0;
|
||
}
|
||
|
||
// 以下函数保持不变
|
||
YOLO_API void free_memory(Detection* detections) {
|
||
delete[] detections;
|
||
}
|
||
|
||
YOLO_API void draw_and_encode_image(
|
||
unsigned char* image_bytes,
|
||
int image_width,
|
||
int image_height,
|
||
const Detection* detections,
|
||
int detections_count,
|
||
const char** class_names,
|
||
int class_names_count,
|
||
unsigned char** out_image_bytes,
|
||
int* out_image_size) {
|
||
|
||
cv::Mat image(image_height, image_width, CV_8UC3, image_bytes);
|
||
if (image.empty()) {
|
||
*out_image_bytes = nullptr;
|
||
*out_image_size = 0;
|
||
return;
|
||
}
|
||
|
||
for (int i = 0; i < detections_count; ++i) {
|
||
const auto& d = detections[i];
|
||
cv::rectangle(image, cv::Rect(d.x, d.y, d.width, d.height), cv::Scalar(0, 255, 0), 2);
|
||
std::string label = "Unknown";
|
||
if (d.class_id >= 0 && d.class_id < class_names_count) {
|
||
label = class_names[d.class_id];
|
||
}
|
||
label += " " + std::to_string(d.score).substr(0, 4);
|
||
cv::putText(image, label, cv::Point(d.x, d.y - 10), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);
|
||
}
|
||
|
||
std::vector<unsigned char> buf;
|
||
cv::imencode(".jpg", image, buf);
|
||
*out_image_size = static_cast<int>(buf.size());
|
||
*out_image_bytes = new unsigned char[*out_image_size];
|
||
std::copy(buf.begin(), buf.end(), *out_image_bytes);
|
||
}
|
||
|
||
YOLO_API void free_image_memory(unsigned char* image_bytes) {
|
||
delete[] image_bytes;
|
||
}
|
||
} |