118 lines
3.5 KiB
C++
118 lines
3.5 KiB
C++
#include <iostream>
|
|
#include <opencv2/opencv.hpp>
|
|
#include <vector> //
|
|
#include "../Yolo11_ONNX/Yolo_ONNX.h" //
|
|
|
|
int main()
|
|
{
|
|
const wchar_t* model_path = L"D:/dev/models/best_fixed_dim.onnx";
|
|
|
|
//
|
|
//
|
|
std::vector<std::string> image_paths = {
|
|
|
|
"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 = 640;
|
|
const int input_height = 640;
|
|
|
|
const char* class_names[] = { "qd", "fl", "zw" };
|
|
int class_count = sizeof(class_names) / sizeof(class_names[0]);
|
|
|
|
|
|
std::cout << "Loading model... (This happens only once)" << std::endl;
|
|
void* detector_handle = create_detector(model_path, input_width, input_height);
|
|
|
|
if (detector_handle == nullptr) {
|
|
std::cerr << "Error: Failed to create detector." << std::endl;
|
|
return -1;
|
|
}
|
|
std::cout << "Model loaded successfully." << 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);
|
|
|
|
if (image.empty()) {
|
|
std::cerr << "Error: Could not read image " << image_path << std::endl;
|
|
continue; //
|
|
}
|
|
|
|
//
|
|
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;
|
|
|
|
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<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 = "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;
|
|
} |