更新为C++ dnncuda版本单张图片检测案例

test/yolov5.cpp

@@ -10,30 +10,11 @@ using std::endl;
 using std::string;
 using std::vector;
 
+static const vector<string> class_name = { "cat", "chicken", "cow", "dog", "fox", "goat", "horse", "person", "racoon", "skunk" };
 
-static const vector<string> class_name = {"cat", "chicken", "cow", "dog", "fox", "goat", "horse", "person", "racoon", "skunk"};
-
-// void print_result(const Mat &result, float conf = 0.7, int len_data = 15)
-// {
-//     float *pdata = (float *)result.data;
-//     for (int i = 0; i < result.total() / len_data; i++)
-//     {
-//         if (pdata[4] > conf)
-//         {
-//             for (int j = 0; j < len_data; j++)
-//             {
-//                 cout << pdata[j] << " ";
-//             }
-//             cout << endl;
-//         }
-//         pdata += len_data;
-//     }
-//     return;
-// }
-
-vector<vector<float>> get_info(const Mat &result, float conf = 0.7, int len_data = 15)
+vector<vector<float>> get_info(const Mat& result, float conf = 0.7, int len_data = 15)
 {
-    float *pdata = (float *)result.data;
+    float* pdata = (float*)result.data;
     vector<vector<float>> info;
     for (int i = 0; i < result.total() / len_data; i++)
     {
@@ -53,7 +34,7 @@ vector<vector<float>> get_info(const Mat &result, float conf = 0.7, int len_data
     return info;
 }
 
-void info_simplify(vector<vector<float>> &info)
+void info_simplify(vector<vector<float>>& info)
 {
     for (auto i = 0; i < info.size(); i++)
     {
@@ -70,7 +51,7 @@ void info_simplify(vector<vector<float>> &info)
     }
 }
 
-vector<vector<vector<float>>> split_info(vector<vector<float>> &info)
+vector<vector<vector<float>>> split_info(vector<vector<float>>& info)
 {
     vector<vector<vector<float>>> info_split;
     vector<int> class_id;
@@ -87,7 +68,7 @@ vector<vector<vector<float>>> split_info(vector<vector<float>> &info)
     return info_split;
 }
 
-void nms(vector<vector<float>> &info, float iou = 0.4)
+void nms(vector<vector<float>>& info, float iou = 0.4)
 {
     int counter = 0;
     vector<vector<float>> return_info;
@@ -99,7 +80,7 @@ void nms(vector<vector<float>> &info, float iou = 0.4)
         float y1 = 0;
         float y2 = 0;
         std::sort(info.begin(), info.end(), [](vector<float> p1, vector<float> p2)
-                  { return p1[4] > p2[4]; });
+            { return p1[4] > p2[4]; });
         for (auto i = 0; i < info.size(); i++)
         {
             if (i < counter)
@@ -139,63 +120,51 @@ void nms(vector<vector<float>> &info, float iou = 0.4)
     }
 }
 
-// void print_info(const vector<vector<float>> &info)
-// {
-//     for (auto i = 0; i < info.size(); i++)
-//     {
-//         for (auto j = 0; j < info[i].size(); j++)
-//         {
-//             cout << info[i][j] << " ";
-//         }
-//         cout << endl;
-//     }
-// }
-
-void draw_box(Mat &img, const vector<vector<float>> &info)
+void draw_box(Mat& img, const vector<vector<float>>& info)
 {
     for (int i = 0; i < info.size(); i++)
     {
-        cv::rectangle(img, cv::Point(info[i][0], info[i][1]), cv::Point(info[i][2], info[i][3]), cv::Scalar(0, 255, 0));
-        string label;
-        label += class_name[info[i][5]];
-        label += "  ";
-        std::stringstream oss;
-        oss << info[i][4];
-        label += oss.str();
-        cv::putText(img, label, cv::Point(info[i][0], info[i][1]), 1, 2, cv::Scalar(0, 255, 0), 2);
-
+        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()
 {
+    // Load the network and set the backend to CUDA
     cv::dnn::Net net = cv::dnn::readNetFromONNX("best.onnx");
+
+    // Set the DNN backend to CUDA and target to CUDA
+    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
+    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
+
     Mat img = cv::imread("fox.jpg");
     cv::resize(img, img, cv::Size(640, 640));
     Mat blob = cv::dnn::blobFromImage(img, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(), true);
     net.setInput(blob);
     vector<Mat> netoutput;
-    vector<string> out_name = {"output"};
+    vector<string> out_name = { "output" };
     net.forward(netoutput, out_name);
     Mat result = netoutput[0];
-    // print_result(result);
     vector<vector<float>> info = get_info(result);
     info_simplify(info);
     vector<vector<vector<float>>> info_split = split_info(info);
-    // cout << " split info" << endl;
-    // print_info(info_split[0]);
-    // cout << info.size() << " " << info[0].size() << endl;
 
-    for(auto i=0; i < info_split.size(); i++)
+    for (auto i = 0; i < info_split.size(); i++)
     {
         nms(info_split[i]);
         draw_box(img, info_split[i]);
     }
-
-    // nms(info_split[0]);
-    // cout << "nms" << endl;
-    // print_info(info_split[0]);
-    // draw_box(img, info_split[0]);
     cv::imshow("test", img);
     cv::waitKey(0);
     return 0;