106 lines
3.3 KiB
C++
106 lines
3.3 KiB
C++
|
#include <iostream>
|
|||
|
|
#include <opencv2/opencv.hpp>
|
|||
|
|
#include "../Yolo11_ONNX/Yolo_ONNX.h" // 包含DLL的头文件
|
|||
|
|
|
|||
|
|
int main()
|
|||
|
|
{
|
|||
|
|
// --- 1. 准备输入和参数 ---
|
|||
|
|
const wchar_t* model_path = L"D:/dev/models/best.onnx";
|
|||
|
|
cv::Mat image = cv::imread("D:/dev/dataset/003.jpg");
|
|||
|
|
|
|||
|
|
if (image.empty()) {
|
|||
|
|
std::cerr << "Error: Could not read the image." << std::endl;
|
|||
|
|
return -1;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 【新增】将所有可配置参数定义在这里，方便修改
|
|||
|
|
const float conf_threshold = 0.9f;
|
|||
|
|
const float iou_threshold = 0.45f;
|
|||
|
|
const int input_width = 2016;
|
|||
|
|
const int input_height = 1536;
|
|||
|
|
|
|||
|
|
const char* class_names[] = { /* ... 您的类别列表 ... */ "tiaojuan", "zhujiesi", "yulingwen" }; // 示例
|
|||
|
|
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;
|
|||
|
|
|
|||
|
|
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);
|
|||
|
|
return -1;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
std::cout << "Detection successful. Found " << detections_count << " objects." << std::endl;
|
|||
|
|
|
|||
|
|
// --- 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<unsigned char> 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;
|
|||
|
|
}
|
|||
|
|
else {
|
|||
|
|
std::cerr << "Error: Failed to save the annotated image." << std::endl;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
// --- 5. 释放内存 ---
|
|||
|
|
std::cout << "Freeing memory..." << std::endl;
|
|||
|
|
free_memory(detections);
|
|||
|
|
free_image_memory(output_image_bytes);
|
|||
|
|
|
|||
|
|
std::cout << "Done." << std::endl;
|
|||
|
|
|
|||
|
|
return 0;
|
|||
|
|
}
|