|
@@ -0,0 +1,280 @@
|
|
|
+#include <array>
|
|
|
+#include <algorithm>
|
|
|
+#include <iostream>
|
|
|
+#include <opencv2/opencv.hpp>
|
|
|
+#include <onnxruntime_cxx_api.h>
|
|
|
+#include <vector>
|
|
|
+#include <string>
|
|
|
+#include <chrono>
|
|
|
+#include <filesystem>
|
|
|
+
|
|
|
+using cv::Mat;
|
|
|
+using std::cout;
|
|
|
+using std::endl;
|
|
|
+using std::string;
|
|
|
+using std::vector;
|
|
|
+namespace fs = std::filesystem;
|
|
|
+
|
|
|
+static const vector<string> class_name = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
|
|
|
+ "fire hydrant","stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
|
|
|
+ "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
|
|
|
+ "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
|
|
|
+ "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
|
|
|
+ "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
|
|
|
+ "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
|
|
+ "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
|
|
|
+ "hair drier", "toothbrush" };
|
|
|
+
|
|
|
+// 筛选置信度低的
|
|
|
+vector<vector<float>> get_info(const float* pdata, int total, float conf = 0.5, int len_data = 85)
|
|
|
+{
|
|
|
+ vector<vector<float>> info;
|
|
|
+ for (int i = 0; i < total / len_data; i++)
|
|
|
+ {
|
|
|
+ if (pdata[4] >= conf)
|
|
|
+ {
|
|
|
+ vector<float> line_data(pdata, pdata + len_data);
|
|
|
+ info.push_back(line_data);
|
|
|
+ }
|
|
|
+ pdata += len_data;
|
|
|
+ }
|
|
|
+ return info;
|
|
|
+}
|
|
|
+
|
|
|
+void info_simplify(vector<vector<float>>& info)
|
|
|
+{
|
|
|
+ for (int i = 0; i < info.size(); i++)
|
|
|
+ {
|
|
|
+ // 找出类别的索引
|
|
|
+ auto max_pos = std::max_element(info[i].cbegin() + 5, info[i].cend());
|
|
|
+ int class_id = std::distance(info[i].cbegin() + 5, max_pos);
|
|
|
+
|
|
|
+ // 仅保留类别为 "person" 的检测
|
|
|
+ if (class_id == 0)
|
|
|
+ {
|
|
|
+ info[i][5] = class_id;
|
|
|
+ info[i].resize(6);
|
|
|
+
|
|
|
+ float x = info[i][0];
|
|
|
+ float y = info[i][1];
|
|
|
+ float w = info[i][2];
|
|
|
+ float h = info[i][3];
|
|
|
+ info[i][0] = x - w / 2.0;
|
|
|
+ info[i][1] = y - h / 2.0;
|
|
|
+ info[i][2] = x + w / 2.0;
|
|
|
+ info[i][3] = y + h / 2.0;
|
|
|
+ }
|
|
|
+ else
|
|
|
+ {
|
|
|
+ info.erase(info.begin() + i);
|
|
|
+ i--;
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+void nms(vector<vector<float>>& info, float iou = 0.4)
|
|
|
+{
|
|
|
+ int counter = 0;
|
|
|
+ vector<vector<float>> return_info;
|
|
|
+ while (counter < info.size())
|
|
|
+ {
|
|
|
+ return_info.clear();
|
|
|
+ float x1 = 0;
|
|
|
+ float x2 = 0;
|
|
|
+ float y1 = 0;
|
|
|
+ float y2 = 0;
|
|
|
+ // 按照置信度排序
|
|
|
+ std::sort(info.begin(), info.end(), [](vector<float> p1, vector<float> p2) {
|
|
|
+ return p1[4] > p2[4];
|
|
|
+ });
|
|
|
+ for (auto i = 0; i < info.size(); i++)
|
|
|
+ {
|
|
|
+ if (i < counter)
|
|
|
+ {
|
|
|
+ return_info.push_back(info[i]);
|
|
|
+ continue;
|
|
|
+ }
|
|
|
+ if (i == counter)
|
|
|
+ {
|
|
|
+ x1 = info[i][0];
|
|
|
+ y1 = info[i][1];
|
|
|
+ x2 = info[i][2];
|
|
|
+ y2 = info[i][3];
|
|
|
+ return_info.push_back(info[i]);
|
|
|
+ continue;
|
|
|
+ }
|
|
|
+ if (info[i][0] > x2 || info[i][2] < x1 || info[i][1] > y2 || info[i][3] < y1)
|
|
|
+ {
|
|
|
+ return_info.push_back(info[i]);
|
|
|
+ }
|
|
|
+ else
|
|
|
+ {
|
|
|
+ float over_x1 = std::max(x1, info[i][0]);
|
|
|
+ float over_y1 = std::max(y1, info[i][1]);
|
|
|
+ float over_x2 = std::min(x2, info[i][2]);
|
|
|
+ float over_y2 = std::min(y2, info[i][3]);
|
|
|
+ float s_over = (over_x2 - over_x1) * (over_y2 - over_y1);
|
|
|
+ float s_total = (x2 - x1) * (y2 - y1) + (info[i][0] - info[i][2]) * (info[i][1] - info[i][3]) - s_over;
|
|
|
+ if (s_over / s_total < iou)
|
|
|
+ {
|
|
|
+ return_info.push_back(info[i]);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ info = return_info;
|
|
|
+ counter += 1;
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+void draw_box(Mat& img, const vector<vector<float>>& info)
|
|
|
+{
|
|
|
+ for (int i = 0; i < info.size(); i++)
|
|
|
+ {
|
|
|
+ if (static_cast<int>(info[i][5]) == 0)
|
|
|
+ {
|
|
|
+ cv::Point topLeft(info[i][0], info[i][1]);
|
|
|
+ cv::Point bottomRight(info[i][2], info[i][3]);
|
|
|
+ int thickness = 2;
|
|
|
+ cv::Scalar color(0, 255, 0);
|
|
|
+ int lineType = cv::LINE_8;
|
|
|
+ const int cornerRadius = 5;
|
|
|
+ cv::rectangle(img, topLeft, bottomRight, color, thickness, lineType);
|
|
|
+ string label = class_name[0] + " " + std::to_string(info[i][4]); // 仅显示 "person" 标签
|
|
|
+ cv::Size textSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.6, 1, nullptr);
|
|
|
+ cv::Rect textBgRect(topLeft.x, topLeft.y - textSize.height - 5, textSize.width + 10, textSize.height + 5);
|
|
|
+ cv::rectangle(img, textBgRect, color, cv::FILLED);
|
|
|
+ cv::putText(img, label, cv::Point(topLeft.x + 5, topLeft.y - 5), cv::FONT_HERSHEY_SIMPLEX, 0.6, CV_RGB(255, 255, 255), 2);
|
|
|
+ }
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+int main()
|
|
|
+{
|
|
|
+ // 定义ONNX模型路径
|
|
|
+ const wchar_t* model_path = L"D://Thework//testrelease//yolov5s.onnx";
|
|
|
+
|
|
|
+ // 初始化 ONNX 运行环境和内存信息
|
|
|
+ Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime");
|
|
|
+ auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
|
|
|
+
|
|
|
+ // 配置会话选项
|
|
|
+ Ort::SessionOptions session_options;
|
|
|
+ session_options.SetIntraOpNumThreads(5);
|
|
|
+ session_options.SetGraphOptimizationLevel(ORT_ENABLE_ALL);
|
|
|
+
|
|
|
+ // 定义输入输出的名称
|
|
|
+ const char* input_names[] = { "images" }; // 根据你的模型输入名调整
|
|
|
+ const char* output_names[] = { "output" }; // 根据你的模型输出名调整
|
|
|
+
|
|
|
+ try {
|
|
|
+ // 创建 ONNX 会话
|
|
|
+ Ort::Session session(env, model_path, session_options);
|
|
|
+
|
|
|
+ // 打开摄像头
|
|
|
+ //cv::VideoCapture cap(0);
|
|
|
+ //if (!cap.isOpened())
|
|
|
+ //{
|
|
|
+ // cout << "Error: Cannot open the camera" << endl;
|
|
|
+ // return -1;
|
|
|
+ //}
|
|
|
+ cv::VideoCapture cap(7, cv::CAP_V4L2);
|
|
|
+ if (!cap.isOpened())
|
|
|
+ {
|
|
|
+ cout << "Error: Cannot open the camera" << endl;
|
|
|
+ return -1;
|
|
|
+ }
|
|
|
+
|
|
|
+ cap.set(cv::CAP_PROP_FRAME_WIDTH, 1280);
|
|
|
+ cap.set(cv::CAP_PROP_FRAME_HEIGHT, 720);
|
|
|
+
|
|
|
+ // 创建保存图像的文件夹
|
|
|
+ std::filesystem::path save_dir = "saved_images";
|
|
|
+ std::filesystem::create_directory(save_dir);
|
|
|
+
|
|
|
+ int image_count = 0; // 用于记录已保存的图像数量
|
|
|
+ auto start_time = std::chrono::steady_clock::now();
|
|
|
+
|
|
|
+ while (image_count < 5)
|
|
|
+ {
|
|
|
+ cv::Mat img;
|
|
|
+ cap >> img;
|
|
|
+ if (img.empty()) break;
|
|
|
+
|
|
|
+ double start = cv::getTickCount();
|
|
|
+
|
|
|
+ // 获取原始图像尺寸
|
|
|
+ cv::Size originalSize = img.size();
|
|
|
+
|
|
|
+ // 将图像调整为模型要求的尺寸
|
|
|
+ cv::Mat resized_img;
|
|
|
+ cv::resize(img, resized_img, cv::Size(640, 640));
|
|
|
+
|
|
|
+ cv::Mat blob = cv::dnn::blobFromImage(resized_img, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(), true);
|
|
|
+
|
|
|
+ // 定义输入Tensor
|
|
|
+ std::array<int64_t, 4> input_shape = { 1, 3, 640, 640 };
|
|
|
+ Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_shape.data(), input_shape.size());
|
|
|
+
|
|
|
+ // 定义输出Tensor的形状
|
|
|
+ std::array<int64_t, 3> output_shape{ 1, 25200, 85 }; // 根据实际模型输出调整
|
|
|
+ std::vector<float> output_tensor_values(1 * 25200 * 85); // 根据实际模型输出调整
|
|
|
+ Ort::Value output_tensor = Ort::Value::CreateTensor<float>(memory_info, output_tensor_values.data(), output_tensor_values.size(), output_shape.data(), output_shape.size());
|
|
|
+
|
|
|
+ // 推理
|
|
|
+ session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, &output_tensor, 1);
|
|
|
+
|
|
|
+ // 处理输出
|
|
|
+ float* output_data = output_tensor.GetTensorMutableData<float>();
|
|
|
+ int output_size = 25200 * 85; // 根据实际模型输出调整
|
|
|
+
|
|
|
+ vector<vector<float>> info = get_info(output_data, output_size, 0.5);
|
|
|
+ info_simplify(info);
|
|
|
+ nms(info);
|
|
|
+
|
|
|
+ // 将检测框重新映射回原始图像尺寸
|
|
|
+ for (auto& box : info)
|
|
|
+ {
|
|
|
+ box[0] *= (float)originalSize.width / 640.0;
|
|
|
+ box[1] *= (float)originalSize.height / 640.0;
|
|
|
+ box[2] *= (float)originalSize.width / 640.0;
|
|
|
+ box[3] *= (float)originalSize.height / 640.0;
|
|
|
+ }
|
|
|
+
|
|
|
+ // 在原始图像上绘制检测框
|
|
|
+ draw_box(img, info);
|
|
|
+
|
|
|
+ double end = cv::getTickCount();
|
|
|
+ double timeSec = (end - start) / cv::getTickFrequency();
|
|
|
+ cout << "Frame time: " << timeSec << " seconds" << endl;
|
|
|
+
|
|
|
+ // 每隔3秒保存一张图像
|
|
|
+ auto now = std::chrono::steady_clock::now();
|
|
|
+ if (std::chrono::duration_cast<std::chrono::seconds>(now - start_time).count() >= 3)
|
|
|
+ {
|
|
|
+ string filename = (save_dir / ("image_" + std::to_string(image_count) + ".jpg")).string();
|
|
|
+ cv::imwrite(filename, img);
|
|
|
+ cout << "Saved " << filename << endl;
|
|
|
+ start_time = now; // 重置计时器
|
|
|
+ image_count++; // 增加计数器
|
|
|
+ }
|
|
|
+
|
|
|
+ // 等待1毫秒
|
|
|
+ std::this_thread::sleep_for(std::chrono::milliseconds(1));
|
|
|
+ }
|
|
|
+
|
|
|
+ // 释放摄像头
|
|
|
+ cap.release();
|
|
|
+ cv::destroyAllWindows();
|
|
|
+ }
|
|
|
+ catch (const Ort::Exception& e) {
|
|
|
+ std::cerr << "ONNX Runtime 异常: " << e.what() << std::endl;
|
|
|
+ }
|
|
|
+ catch (const std::exception& e) {
|
|
|
+ std::cerr << "标准异常: " << e.what() << std::endl;
|
|
|
+ }
|
|
|
+ catch (...) {
|
|
|
+ std::cerr << "未知异常." << std::endl;
|
|
|
+ }
|
|
|
+
|
|
|
+ return 0;
|
|
|
+}
|