|
@@ -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;
|
|
|
+}
|