C++运行YOLO进行目标检测

使用像Ultralytics这样的用户友好库,用Python运行目标检测模型已经变得非常简单,但是如何用C++运行YOLO模型呢?

C++运行YOLO进行目标检测

本文是关于如何在CPU上运行YOLOv5模型,而不是GPU。要在GPU上运行模型,需要安装CUDA、CUDNN等,这可能会让人感到困惑且耗时。我将在未来写另一篇文章介绍如何在支持CUDA的情况下运行YOLO模型。

目前,你只需要安装OpenCV库。如果你还没有安装它,可以从这个链接安装。

那么,让我们开始吧。我会一步一步地讲解,完整的代码可以在页面底部查看。

1、导出YOLOv5模型为ONNX模型

创建一个新的文件夹并命名为你想要的名称。打开终端并将yolov5仓库克隆到该文件夹中。我们将使用此仓库导出模型为onnx格式。

git clone https://github.com/ultralytics/yolov5

我将使用yolov5s.pt模型,但你可以使用自定义的YOLOv5模型,过程不会改变。你可以从这个链接下载预训练模型,或者使用你自己的模型。如果你不知道如何训练自定义YOLO模型,不用担心,我已经有一篇文章介绍了这个内容,你可以查看它(训练自定义YOLO模型

预训练YOLOv5模型

现在让我们将模型导出为onnx。有不同参数,你可以参考下面的图片。你可以编辑“export.py”文件(yolov5/export.py),或者像我在这里一样手动更改参数。对于自定义模型,你需要将权重更改为你的自定义模型权重(your_model.pt文件)。

python yolov5/export.py --weights yolov5s.pt --img 640 --include onnx --opset 12

注意:你需要将opset设置为12,否则可能会报错。这是一个常见的问题,你可以在GitHub上查找此错误。

parameers/ export.py文件

2、创建TXT文件存储YOLO模型标签

这一步很简单,你只需要创建一个txt文件来存储标签。如果你使用的是预训练的YOLO模型,可以从此链接直接下载txt文件。
如果你有一个自定义模型,则需要创建一个新的txt文件并在其中写入你的标签。你可以给这个文件命名任何你想要的名字,这并不重要。

coco-classes.txt

3、创建CMakeLists.txt文件

现在,让我们创建一个CMakeLists.txt文件。当你使用CMake编译C++程序时需要这个文件。如果你按照我提供的链接安装了OpenCV,你应该已经安装了CMake。

cmake_minimum_required(VERSION 3.10)  
project(cpp-yolo-detection) # 你的文件夹名称
  
# 查找OpenCV  
set(OpenCV_DIR C:/Libraries/opencv/build) # OpenCV路径  
find_package(OpenCV REQUIRED)  
  
add_executable(object-detection object-detection.cpp) # 你的文件名  
  
# 链接OpenCV库  
target_link_libraries(object-detection ${OpenCV_LIBS})

4、代码

这是最后一步。我使用了这个仓库中的代码,但修改了一些部分,并添加了注释以便更好地理解。

#include <fstream>  
#include <opencv2/opencv.hpp>  
  
  
// 从coco-classes.txt文件加载标签  
std::vector<std::string> load_class_list()  
{  
    std::vector<std::string> class_list;  
    // 更改此txt文件为你包含标签的txt文件  
    std::ifstream ifs("C:/Users/sirom/Desktop/cpp-ultralytics/coco-classes.txt");  
    std::string line;  
    while (getline(ifs, line))  
    {  
        class_list.push_back(line);  
    }  
    return class_list;  
}  
  
// 模型   
void load_net(cv::dnn::Net &net)  
{     
    // 更改此路径为你模型的路径   
    auto result = cv::dnn::readNet("C:/Users/sirom/Desktop/cpp-ultralytics/yolov5s.onnx");  
  
    std::cout << "在CPU上运行/n";  
    result.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);  
    result.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);  
   
    net = result;  
}  
  
const std::vector<cv::Scalar> colors = {cv::Scalar(255, 255, 0), cv::Scalar(0, 255, 0), cv::Scalar(0, 255, 255), cv::Scalar(255, 0, 0)};  
  
// 你可以更改这些参数以获得更好的结果  
const float INPUT_WIDTH = 640.0;  
const float INPUT_HEIGHT = 640.0;  
const float SCORE_THRESHOLD = 0.5;  
const float NMS_THRESHOLD = 0.5;  
const float CONFIDENCE_THRESHOLD = 0.5;  
  
struct Detection  
{  
    int class_id;  
    float confidence;  
    cv::Rect box;  
};  
  
// yolov5格式  
cv::Mat format_yolov5(const cv::Mat &source) {  
    int col = source.cols;  
    int row = source.rows;  
    int _max = MAX(col, row);  
    cv::Mat result = cv::Mat::zeros(_max, _max, CV_8UC3);  
    source.copyTo(result(cv::Rect(0, 0, col, row)));  
    return result;  
}  
  
// 检测函数  
void detect(cv::Mat &image, cv::dnn::Net &net, std::vector<Detection> &output, const std::vector<std::string> &className) {  
    cv::Mat blob;  
  
    // 将输入图像格式化以满足模型输入要求  
    auto input_image = format_yolov5(image);  
      
    // 将图像转换为blob并将其作为网络的输入  
    cv::dnn::blobFromImage(input_image, blob, 1./255., cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::Scalar(), true, false);  
    net.setInput(blob);  
    std::vector<cv::Mat> outputs;  
    net.forward(outputs, net.getUnconnectedOutLayersNames());  
  
    // 缩放因子用于将边界框映射回原始图像大小  
    float x_factor = input_image.cols / INPUT_WIDTH;  
    float y_factor = input_image.rows / INPUT_HEIGHT;  
      
    float *data = (float *)outputs[0].data;  
  
    const int dimensions = 85;  
    const int rows = 25200;  
      
    std::vector<int> class_ids; // 存储检测的类别ID  
    std::vector<float> confidences; // 存储检测的置信度  
    std::vector<cv::Rect> boxes;   // 存储边界框  
  
   // 循环处理所有行以处理预测  
    for (int i = 0; i < rows; ++i) {  
  
        // 获取当前检测的置信度  
        float confidence = data[4];  
  
        // 只处理置信度高于阈值的检测  
        if (confidence >= CONFIDENCE_THRESHOLD) {  
              
            // 获取类别分数并找到最高分的类别  
            float * classes_scores = data + 5;  
            cv::Mat scores(1, className.size(), CV_32FC1, classes_scores);  
            cv::Point class_id;  
            double max_class_score;  
            minMaxLoc(scores, 0, &max_class_score, 0, &class_id);  
  
            // 如果类别分数高于阈值,则存储检测  
            if (max_class_score > SCORE_THRESHOLD) {  
  
                confidences.push_back(confidence);  
                class_ids.push_back(class_id.x);  
  
                // 计算边界框坐标  
                float x = data[0];  
                float y = data[1];  
                float w = data[2];  
                float h = data[3];  
                int left = int((x - 0.5 * w) * x_factor);  
                int top = int((y - 0.5 * h) * y_factor);  
                int width = int(w * x_factor);  
                int height = int(h * y_factor);  
                boxes.push_back(cv::Rect(left, top, width, height));  
            }  
        }  
  
        data += 85;  
    }  
  
    // 应用非最大抑制  
    std::vector<int> nms_result;  
    cv::dnn::NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, nms_result);  
  
    // 绘制NMS过滤后的边界框并推送结果  
    for (int i = 0; i < nms_result.size(); i++) {  
        int idx = nms_result[i];  
  
        // 只推送过滤后的检测  
        Detection result;  
        result.class_id = class_ids[idx];  
        result.confidence = confidences[idx];  
        result.box = boxes[idx];  
        output.push_back(result);  
  
        // 绘制最终的NMS边界框和标签  
        cv::rectangle(image, boxes[idx], cv::Scalar(0, 255, 0), 3);  
        std::string label = className[class_ids[idx]];  
        cv::putText(image, label, cv::Point(boxes[idx].x, boxes[idx].y - 5), cv::FONT_HERSHEY_SIMPLEX, 2, cv::Scalar(255, 255, 255), 2);  
    }  
}  
  
  
int main(int argc, char **argv)  
{     
    // 加载类别列表   
    std::vector<std::string> class_list = load_class_list();  
  
    // 加载输入图像  
    std::string image_path = cv::samples::findFile("C:/Users/sirom/Desktop/cpp-ultralytics/test2.jpg");  
    cv::Mat frame = cv::imread(image_path, cv::IMREAD_COLOR);  
  
    // 加载模型  
    cv::dnn::Net net;  
    load_net(net);  
  
    // 向量用于存储检测结果  
    std::vector<Detection> output;  
    // 在输入图像上运行检测  
    detect(frame, net, output, class_list);  
  
    // 将结果保存到文件  
    cv::imwrite("C:/Users/sirom/Desktop/cpp-ultralytics/result.jpg", frame);  
  
    while (true)  
    {         
        // 显示图像  
        cv::imshow("image",frame);  
  
        // 如果按任意键则退出循环  
        if (cv::waitKey(1) != -1)  
        {  
            std::cout << "由用户完成\n";  
            break;  
        }  
    }  
  
    std::cout << "处理完成。图像已保存 /n";  
    return 0;  
}

5、编译并运行代码

mkdir build  
cd build   
cmake ..  
cmake --build .  
.\Debug\object-detection.exe

原文链接:How to Run YOLO Models in C++ for Object Detection

汇智网翻译整理,转载请标明出处