diff --git a/TestApp/TestApp.cpp b/TestApp/TestApp.cpp index 58cd1e4..386cbc2 100644 --- a/TestApp/TestApp.cpp +++ b/TestApp/TestApp.cpp @@ -1,106 +1,118 @@ #include #include -#include "../Yolo11_ONNX/Yolo_ONNX.h" // 包含DLL的头文件 +#include // +#include "../Yolo11_ONNX/Yolo_ONNX.h" // int main() { - // --- 1. 准备输入和参数 --- - const wchar_t* model_path = L"D:/dev/models/best.onnx"; - cv::Mat image = cv::imread("D:/dev/dataset/003.jpg"); + const wchar_t* model_path = L"D:/dev/models/best_fixed_dim.onnx"; - if (image.empty()) { - std::cerr << "Error: Could not read the image." << std::endl; - return -1; - } + // + // + std::vector image_paths = { - // 【新增】将所有可配置参数定义在这里，方便修改 - const float conf_threshold = 0.9f; + "D:/dev/dataset/qd/1.jpg", + "D:/dev/dataset/qd/2.jpg", + "D:/dev/dataset/qd/3.jpg", + "D:/dev/dataset/qd/4.jpg", + "D:/dev/dataset/qd/5.jpg" + }; + + const float conf_threshold = 0.15f; const float iou_threshold = 0.45f; - const int input_width = 2016; - const int input_height = 1536; + const int input_width = 640; + const int input_height = 640; - const char* class_names[] = { /* ... 您的类别列表 ... */ "tiaojuan", "zhujiesi", "yulingwen" }; // 示例 + const char* class_names[] = { "qd", "fl", "zw" }; int class_count = sizeof(class_names) / sizeof(class_names[0]); - // --- 2. 调用更新后的DLL函数 --- - Detection* detections = nullptr; - int detections_count = 0; - std::cout << "Performing detection with conf=" << conf_threshold << ", iou=" << iou_threshold << ", size=" << input_width << "x" << input_height << std::endl; + std::cout << "Loading model... (This happens only once)" << std::endl; + void* detector_handle = create_detector(model_path, input_width, input_height); - int result = perform_detection( - model_path, - image.data, - image.cols, - image.rows, - &detections, - &detections_count, - class_names, - class_count, - conf_threshold, - iou_threshold, - input_width, - input_height - ); - - if (result != 0) { - std::cerr << "Detection failed with code: " << result << std::endl; - free_memory(detections); + if (detector_handle == nullptr) { + std::cerr << "Error: Failed to create detector." << std::endl; return -1; } + std::cout << "Model loaded successfully." << std::endl; - std::cout << "Detection successful. Found " << detections_count << " objects." << std::endl; + for (const auto& image_path : image_paths) + { + std::cout << "\n--- Processing image: " << image_path << " ---" << std::endl; + cv::Mat image = cv::imread(image_path); - // --- 3. 打印检测结果 --- - for (int i = 0; i < detections_count; ++i) { - const auto& d = detections[i]; - std::cout << " - Class: " << class_names[d.class_id] - << ", Score: " << d.score - << ", Box: [" << d.x << ", " << d.y << ", " << d.width << ", " << d.height << "]" << std::endl; - } - - // --- 4. 调用DLL函数绘制结果并显示 --- - unsigned char* output_image_bytes = nullptr; - int output_image_size = 0; - - draw_and_encode_image( - image.data, - image.cols, - image.rows, - detections, - detections_count, - class_names, - class_count, - &output_image_bytes, - &output_image_size - ); - - if (output_image_bytes && output_image_size > 0) { - std::vector buffer(output_image_bytes, output_image_bytes + output_image_size); - cv::Mat result_image = cv::imdecode(buffer, cv::IMREAD_COLOR); - - // 定义要保存的文件名 - std::string output_filename = "detection_result.jpg"; - - // 使用 OpenCV 的 imwrite 函数将图片保存到硬盘 - bool success = cv::imwrite(output_filename, result_image); - - // 检查是否保存成功并打印提示信息 - if (success) { - std::cout << "Annotated image successfully saved to: " << output_filename << std::endl; + if (image.empty()) { + std::cerr << "Error: Could not read image " << image_path << std::endl; + continue; // } - else { - std::cerr << "Error: Failed to save the annotated image." << std::endl; + + // + Detection* detections = nullptr; + int detections_count = 0; + + // + int result = perform_detection_on_session( + detector_handle, // + image.data, + image.cols, + image.rows, + &detections, + &detections_count, + conf_threshold, + iou_threshold + ); + + if (result != 0) { + std::cerr << "Detection failed with code: " << result << std::endl; + free_memory(detections); // + continue; } - } + std::cout << "Detection successful. Found " << detections_count << " objects." << std::endl; - // --- 5. 释放内存 --- - std::cout << "Freeing memory..." << std::endl; - free_memory(detections); - free_image_memory(output_image_bytes); + for (int i = 0; i < detections_count; ++i) { + const auto& d = detections[i]; + std::cout << " - Class: " << class_names[d.class_id] + << ", Score: " << d.score + << ", Box: [" << d.x << ", " << d.y << ", " << d.width << ", " << d.height << "]" << std::endl; + } + + unsigned char* output_image_bytes = nullptr; + int output_image_size = 0; + + draw_and_encode_image( + image.data, + image.cols, + image.rows, + detections, + detections_count, + class_names, + class_count, + &output_image_bytes, + &output_image_size + ); + + if (output_image_bytes && output_image_size > 0) { + std::vector buffer(output_image_bytes, output_image_bytes + output_image_size); + cv::Mat result_image = cv::imdecode(buffer, cv::IMREAD_COLOR); + + // + std::string output_filename = "result_" + image_path.substr(image_path.find_last_of('/') + 1); + cv::imwrite(output_filename, result_image); + std::cout << "Annotated image saved to: " << output_filename << std::endl; + } + + // --- 5. [修改] + // + free_memory(detections); + free_image_memory(output_image_bytes); + + } // + + // --- 6. [新增] + std::cout << "\nFreeing detector model..." << std::endl; + free_detector(detector_handle); std::cout << "Done." << std::endl; - return 0; } \ No newline at end of file diff --git a/TestApp/result_001.jpg b/TestApp/result_001.jpg new file mode 100644 index 0000000..e539881 Binary files /dev/null and b/TestApp/result_001.jpg differ diff --git a/TestApp/result_002.jpg b/TestApp/result_002.jpg new file mode 100644 index 0000000..c4d2ecb Binary files /dev/null and b/TestApp/result_002.jpg differ diff --git a/TestApp/result_003.jpg b/TestApp/result_003.jpg new file mode 100644 index 0000000..63ca5b9 Binary files /dev/null and b/TestApp/result_003.jpg differ diff --git a/TestApp/result_1.jpg b/TestApp/result_1.jpg new file mode 100644 index 0000000..c4d2ecb Binary files /dev/null and b/TestApp/result_1.jpg differ diff --git a/TestApp/result_2.jpg b/TestApp/result_2.jpg new file mode 100644 index 0000000..e89a125 Binary files /dev/null and b/TestApp/result_2.jpg differ diff --git a/TestApp/result_3.jpg b/TestApp/result_3.jpg new file mode 100644 index 0000000..24bae57 Binary files /dev/null and b/TestApp/result_3.jpg differ diff --git a/TestApp/result_4.jpg b/TestApp/result_4.jpg new file mode 100644 index 0000000..181b590 Binary files /dev/null and b/TestApp/result_4.jpg differ diff --git a/TestApp/result_5.jpg b/TestApp/result_5.jpg new file mode 100644 index 0000000..955a0ea Binary files /dev/null and b/TestApp/result_5.jpg differ diff --git a/TestApp/result_test001.jpg b/TestApp/result_test001.jpg new file mode 100644 index 0000000..4115705 Binary files /dev/null and b/TestApp/result_test001.jpg differ diff --git a/Yolo11_ONNX/Yolo_ONNX.cpp b/Yolo11_ONNX/Yolo_ONNX.cpp index 5c09afe..71c6946 100644 --- a/Yolo11_ONNX/Yolo_ONNX.cpp +++ b/Yolo11_ONNX/Yolo_ONNX.cpp @@ -1,8 +1,7 @@ -// YoloV8_ONNX.cpp (Ż) - #include "pch.h" #include "Yolo_ONNX.h" +// #include #include #include @@ -10,10 +9,101 @@ #include #include #include +#include // namespace { - // ŻԤ + cv::Mat preprocess(const cv::Mat& img, int target_width, int target_height, int& pad_w, int& pad_h, float& scale); + + std::vector 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); + + class YoloDetector { + public: + // + Ort::Env env; + std::unique_ptr session; + + // + int input_width = 0; + int input_height = 0; + + // + Ort::AllocatorWithDefaultOptions allocator; + std::string input_name_str; + std::string output_name_str; + std::vector input_node_names; + std::vector output_node_names; + + public: + /** + * @brief + */ + YoloDetector(const wchar_t* model_path, int in_width, int in_height) + : env(ORT_LOGGING_LEVEL_WARNING, "YOLOv8-ONNX-GPU"), + input_width(in_width), + input_height(in_height) + { + // + Ort::SessionOptions session_options; + OrtCUDAProviderOptions cuda_options; // + session_options.AppendExecutionProvider_CUDA(cuda_options); + + // + session = std::make_unique(env, model_path, session_options); + + // + // + input_name_str = session->GetInputNameAllocated(0, allocator).get(); + output_name_str = session->GetOutputNameAllocated(0, allocator).get(); + input_node_names.push_back(input_name_str.c_str()); + output_node_names.push_back(output_name_str.c_str()); + } + + /** + * @brief + */ + std::vector detect( + unsigned char* image_bytes, + int image_width, + int image_height, + float conf_threshold, + float iou_threshold) + { + // + cv::Mat image(image_height, image_width, CV_8UC3, image_bytes); + if (image.empty()) { + throw std::runtime_error("Input image is empty."); + } + + // + 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(), true, false); + std::vector input_shape = { 1, 3, (int64_t)input_height, (int64_t)input_width }; + std::cout <<"input shape: " << input_shape[0] << "," << input_shape[1] << "," << input_shape[2] << "," << input_shape[3] << "," << std::endl; + + // + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + Ort::Value input_tensor = Ort::Value::CreateTensor(memory_info, blob.ptr(), 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); + + // + return postprocess(output_tensors[0], scale, pad_w, pad_h, image_width, image_height, conf_threshold, iou_threshold); + } + }; + + + // ======================================================================== + // + // ======================================================================== + + // 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; @@ -21,31 +111,35 @@ namespace { scale = std::min(static_cast(target_width) / w, static_cast(target_height) / h); int new_w = static_cast(w * scale); int new_h = static_cast(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; + pad_w = target_width - new_w; // + pad_h = target_height - new_h; // + + // + int top = pad_h / 2; + int bottom = pad_h - top; + int left = pad_w / 2; + int right = pad_w - left; + // + cv::Mat padded_img; - cv::copyMakeBorder(resized_img, padded_img, 0, pad_h, 0, pad_w, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114)); + // + cv::copyMakeBorder(resized_img, padded_img, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114)); return padded_img; } - // + // std::vector 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(); int num_classes = static_cast(output_shape[1]) - 4; int num_proposals = static_cast(output_shape[2]); - std::vector boxes; std::vector scores; std::vector 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(i); const float* class_scores = proposal + 4; @@ -63,7 +157,7 @@ namespace { float w = proposal[2]; float h = proposal[3]; int left = static_cast((cx - w / 2 - (pad_w / 2.0f)) / scale); - int top = static_cast((cy - h / 2 - (pad_h / 2.0f)) / scale); + int top = static_cast((cy - h / 2 - (pad_w / 2.0f)) / scale); int width = static_cast(w / scale); int height = static_cast(h / scale); left = std::max(0, std::min(left, img_w - 1)); @@ -75,7 +169,6 @@ namespace { class_ids.push_back(class_id); } } - std::vector nms_result; cv::dnn::NMSBoxes(boxes, scores, conf_threshold, iou_threshold, nms_result); std::vector detections; @@ -84,11 +177,92 @@ namespace { } return detections; } -} +} // extern "C" { - // ޸ġǩ£ĸ² + + // ======================================================================== + // + // ======================================================================== + + YOLO_API void* create_detector( + const wchar_t* model_path, + int input_width, + int input_height) + { + try { + // + YoloDetector* detector = new YoloDetector(model_path, input_width, input_height); + // + return static_cast(detector); + } + catch (const Ort::Exception& e) { + std::cerr << "ONNX Runtime Error creating detector: " << e.what() << std::endl; + return nullptr; + } + catch (const std::exception& e) { + std::cerr << "Error creating detector: " << e.what() << std::endl; + return nullptr; + } + } + + YOLO_API void free_detector(void* detector_handle) + { + if (detector_handle) { + YoloDetector* detector = static_cast(detector_handle); + delete detector; + } + } + + YOLO_API int perform_detection_on_session( + void* detector_handle, + unsigned char* image_bytes, + int image_width, + int image_height, + Detection** out_detections, + int* out_detections_count, + float conf_threshold, + float iou_threshold) + { + if (!detector_handle) return -1; // + + // + YoloDetector* detector = static_cast(detector_handle); + + try { + // + std::vector detections = detector->detect( + image_bytes, image_width, image_height, + conf_threshold, iou_threshold + ); + + // + *out_detections_count = static_cast(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; + } + return 0; // + } + catch (const Ort::Exception& e) { + std::cerr << "ONNX Runtime Error during detection: " << e.what() << std::endl; + return -2; + } + catch (const std::exception& e) { + std::cerr << "Error during detection: " << e.what() << std::endl; + return -4; + } + } + + + // ======================================================================== + // + // ======================================================================== + YOLO_API int perform_detection( const wchar_t* model_path, unsigned char* image_bytes, @@ -98,18 +272,17 @@ extern "C" { 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 // ʹôģ߶ + float conf_threshold, + float iou_threshold, + int input_width, + int input_height ) { + // static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "YOLOv8-ONNX-GPU"); static std::unique_ptr 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; @@ -118,13 +291,10 @@ extern "C" { current_model_path = model_path; } - // ޸ġƳӲijߴ磬ʹýӿڴIJ std::vector input_shape = { 1, 3, input_height, input_width }; - Ort::AllocatorWithDefaultOptions allocator; std::string input_name_str = session->GetInputNameAllocated(0, allocator).get(); std::vector input_node_names = { input_name_str.c_str() }; - std::string output_name_str = session->GetOutputNameAllocated(0, allocator).get(); std::vector output_node_names = { output_name_str.c_str() }; @@ -133,17 +303,15 @@ extern "C" { int pad_w, pad_h; float scale; - // ޸ġʹýӿڴIJԤ 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); + cv::dnn::blobFromImage(preprocessed_img, blob, 1 / 255.0, cv::Size(), cv::Scalar(), true, false); auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); Ort::Value input_tensor = Ort::Value::CreateTensor(memory_info, blob.ptr(), 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); - // ޸ġʹýӿڴIJк std::vector detections = postprocess(output_tensors[0], scale, pad_w, pad_h, image_width, image_height, conf_threshold, iou_threshold); *out_detections_count = static_cast(detections.size()); @@ -170,7 +338,7 @@ extern "C" { return 0; } - // ºֲ + // YOLO_API void free_memory(Detection* detections) { delete[] detections; } diff --git a/Yolo11_ONNX/Yolo_ONNX.h b/Yolo11_ONNX/Yolo_ONNX.h index 4de853f..e47956e 100644 --- a/Yolo11_ONNX/Yolo_ONNX.h +++ b/Yolo11_ONNX/Yolo_ONNX.h @@ -1,5 +1,6 @@ #pragma once +// #ifdef YOLO_EXPORTS #define YOLO_API __declspec(dllexport) #else @@ -8,6 +9,7 @@ #include +// struct Detection { int class_id; @@ -19,6 +21,63 @@ struct Detection }; extern "C" { + + // ======================================================================== + // + // ======================================================================== + + /** + * @brief [] һʵ (״̬) + * @details ˺ONNXģͲʼSessionһ + * @param model_path ONNXģ͵ľ· + * @param input_width ģ ( 640) + * @param input_height ģ߶ ( 640) + * @return void* ָʵľʧ򷵻 nullptr + */ + YOLO_API void* create_detector( + const wchar_t* model_path, + int input_width, + int input_height + ); + + /** + * @brief [] ͷ create_detector ļʵ + * @param detector_handle Ҫͷŵľ + */ + YOLO_API void free_detector(void* detector_handle); + + /** + * @brief [] ʹѼصļִм (̰߳ȫ) + * @details + * @param detector_handle create_detector صľ + * @param image_bytes ָͼ(BGRʽ)ָ + * @param image_width ͼ + * @param image_height ͼ߶ + * @param out_detections [] Detectionṹָ + * @param out_detections_count [] ⵽ + * @param conf_threshold Ŷֵ + * @param iou_threshold NMSIOUֵ + * @return int 0ʾɹ + */ + YOLO_API int perform_detection_on_session( + void* detector_handle, + unsigned char* image_bytes, + int image_width, + int image_height, + Detection** out_detections, + int* out_detections_count, + float conf_threshold, + float iou_threshold + ); + + + // ======================================================================== + // + // ======================================================================== + + /** + * @brief [] ԭʼ״̬⺯ (̰߳ȫ) + */ YOLO_API int perform_detection( const wchar_t* model_path, unsigned char* image_bytes, @@ -34,8 +93,14 @@ extern "C" { int input_height // ģ߶ ( 640) ); + /** + * @brief [] ͷ perform_detection* ڴ + */ YOLO_API void free_memory(Detection* detections); + /** + * @brief [] ƺ + */ YOLO_API void draw_and_encode_image( unsigned char* image_bytes, int image_width, @@ -48,5 +113,8 @@ extern "C" { int* out_image_size ); + /** + * @brief [] ͷ draw_and_encode_image ڴ + */ YOLO_API void free_image_memory(unsigned char* image_bytes); } \ No newline at end of file