Эх сурвалжийг харах

yolov5C++版本,onnxruntimeGPU加速,调用系统摄像头。debug没问题,存在内存问题,待修改

Curious 8 сар өмнө
parent
commit
ad6ba95840

+ 239 - 0
test/yolov5C++onnxruntimeGPU.cpp

@@ -0,0 +1,239 @@
+#include <iostream>
+#include <opencv2/opencv.hpp>
+#include <onnxruntime_cxx_api.h>
+#include <vector>
+#include <string>
+#include <algorithm>
+#include <numeric>
+#include <sstream>
+
+using cv::Mat;
+using std::cout;
+using std::endl;
+using std::string;
+using std::vector;
+
+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 Mat& result, float conf = 0.5, int len_data = 85)
+{
+    float* pdata = (float*)result.data;
+    vector<vector<float>> info;
+    for (int i = 0; i < result.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 (auto& entry : info)
+    {
+        entry[5] = std::max_element(entry.begin() + 5, entry.end()) - entry.begin() - 5;
+        entry.resize(6);
+        float x = entry[0];
+        float y = entry[1];
+        float w = entry[2];
+        float h = entry[3];
+        entry[0] = x - w / 2.0;
+        entry[1] = y - h / 2.0;
+        entry[2] = x + w / 2.0;
+        entry[3] = y + h / 2.0;
+    }
+}
+
+vector<vector<vector<float>>> split_info(vector<vector<float>>& info)
+{
+    vector<vector<vector<float>>> info_split;
+    vector<int> class_id;
+
+    for (const auto& entry : info)
+    {
+        int classIndex = static_cast<int>(entry[5]);
+        if (std::find(class_id.begin(), class_id.end(), classIndex) == class_id.end())
+        {
+            class_id.push_back(classIndex);
+            info_split.emplace_back();
+        }
+        info_split[std::find(class_id.begin(), class_id.end(), classIndex) - class_id.begin()].push_back(entry);
+    }
+    return info_split;
+}
+
+void nms(vector<vector<float>>& info, float iou = 0.5)
+{
+    std::sort(info.begin(), info.end(), [](vector<float> p1, vector<float> p2) {
+        return p1[4] > p2[4];
+        });
+    vector<vector<float>> return_info;
+    vector<bool> suppressed(info.size(), false);
+    for (size_t i = 0; i < info.size(); i++)
+    {
+        if (suppressed[i]) continue;
+
+        return_info.push_back(info[i]);
+        float x1 = info[i][0], y1 = info[i][1], x2 = info[i][2], y2 = info[i][3];
+
+        for (size_t j = i + 1; j < info.size(); j++)
+        {
+            if (suppressed[j]) continue;
+
+            float interX1 = std::max(x1, info[j][0]);
+            float interY1 = std::max(y1, info[j][1]);
+            float interX2 = std::min(x2, info[j][2]);
+            float interY2 = std::min(y2, info[j][3]);
+
+            float interArea = std::max(0.0f, interX2 - interX1) * std::max(0.0f, interY2 - interY1);
+            float totalArea = (x2 - x1) * (y2 - y1) + (info[j][2] - info[j][0]) * (info[j][3] - info[j][1]) - interArea;
+            if (interArea / totalArea > iou)
+            {
+                suppressed[j] = true;
+            }
+        }
+    }
+
+    info = std::move(return_info);
+}
+
+void draw_box(Mat& img, const vector<vector<float>>& info)
+{
+    for (int i = 0; i < info.size(); i++)
+    {
+        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[static_cast<int>(info[i][5])] + " " + std::to_string(info[i][4]);
+        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() {
+    try {
+        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntime");
+
+        Ort::SessionOptions session_options;
+        OrtCUDAProviderOptions cuda_options;
+        session_options.AppendExecutionProvider_CUDA(cuda_options);
+
+        std::wstring model_path = L"D:/VisualStudio/test/Project1/Project1/yolov5s.onnx";
+
+        Ort::Session session(env, model_path.c_str(), session_options);
+
+        std::vector<std::string> input_names = { "input" };
+        std::vector<std::string> output_names = { "output" };
+
+        std::vector<const char*> input_names_cstr;
+        std::vector<const char*> output_names_cstr;
+        for (const auto& name : input_names) {
+            input_names_cstr.push_back(name.c_str());
+        }
+        for (const auto& name : output_names) {
+            output_names_cstr.push_back(name.c_str());
+        }
+
+        cv::VideoCapture cap(0);
+        if (!cap.isOpened()) {
+            cout << "Error: Cannot open the camera" << endl;
+            return -1;
+        }
+
+        Ort::MemoryInfo memory_info("Cpu", OrtArenaAllocator, 0, OrtMemTypeDefault);
+
+        while (true) {
+            Mat img;
+            cap >> img;
+            if (img.empty()) break;
+            cv::resize(img, img, cv::Size(640, 640));
+
+            // Convert image to float and normalize
+            Mat blob;
+            img.convertTo(blob, CV_32F, 1.0 / 255.0);
+            std::vector<float> input_tensor_values(blob.total());
+            std::memcpy(input_tensor_values.data(), blob.data, blob.total() * sizeof(float));
+
+            // Create ONNX Runtime tensor from input data
+            std::vector<int64_t> input_node_dims = { 1, 3, 640, 640 };
+            Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
+                memory_info, input_tensor_values.data(), input_tensor_values.size(),
+                input_node_dims.data(), input_node_dims.size());
+
+            // Run inference
+            auto output_tensors = session.Run(Ort::RunOptions{ nullptr },
+                input_names_cstr.data(), &input_tensor, 1,
+                output_names_cstr.data(), 1);
+
+            // Get the result
+            Ort::Value& output_tensor = output_tensors.front();
+            auto output_shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();
+            std::vector<float> output_data(output_tensor.GetTensorTypeAndShapeInfo().GetElementCount());
+
+            // Print output shape for debugging
+            cout << "Output Shape: ";
+            for (const auto& dim : output_shape) {
+                cout << dim << " ";
+            }
+            cout << endl;
+
+            size_t expected_size = std::accumulate(output_shape.begin(), output_shape.end(), 1, std::multiplies<size_t>());
+            if (output_data.size() != expected_size) {
+                std::cerr << "Error: Output data size mismatch." << endl;
+                return -1;
+            }
+
+            std::memcpy(output_data.data(), output_tensor.GetTensorData<float>(),
+                output_data.size() * sizeof(float));
+
+            // Print a portion of the output data for debugging
+            cout << "Output Data (first 10 values): ";
+            for (size_t i = 0; i < std::min(output_data.size(), static_cast<size_t>(10)); ++i) {
+                cout << output_data[i] << " ";
+            }
+            cout << endl;
+
+            Mat result(output_shape[2], output_shape[3], CV_32FC1, output_data.data());
+
+            vector<vector<float>> info = get_info(result);
+            info_simplify(info);
+            vector<vector<vector<float>>> info_split = split_info(info);
+            for (auto& split_info : info_split) {
+                nms(split_info);
+                draw_box(img, split_info);
+            }
+
+            cv::imshow("test", img);
+            if (cv::waitKey(1) == 'q') break;
+        }
+        cap.release();
+        cv::destroyAllWindows();
+    }
+    catch (const Ort::Exception& e) {
+        cout << "ONNX Runtime error: " << e.what() << endl;
+        return -1;
+    }
+    catch (const std::exception& e) {
+        cout << "Standard exception: " << e.what() << endl;
+        return -1;
+    }
+    return 0;
+}