vehicle-demo/yolov5onnxC++/yolov5onnxC++.cpp

#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;
}