基于C++和ONNX Runtime的YOLOv5目标检测实战

1. 前言

在计算机视觉领域,目标检测是一项关键任务,其应用广泛,涵盖了安防监控、自动驾驶、工业检测等众多领域。YOLOv5作为一种先进的目标检测算法,以其速度快、精度高的特点备受关注。本文将详细介绍如何使用C++结合ONNX Runtime推理引擎来部署YOLOv5模型,实现高效的目标检测。

2. ONNX与YOLOv5

2.1 ONNX简介

ONNX(Open Neural Network Exchange)是一种开放的神经网络交换格式,旨在促进不同深度学习框架之间的模型互操作性。它允许开发者在一个框架中训练模型,然后将其转换为ONNX格式,并在另一个支持ONNX的框架或推理引擎中进行部署,大大提高了模型的可移植性和灵活性。

2.2 YOLOv5模型

YOLOv5是YOLO系列目标检测算法的最新版本,其在保持高检测精度的同时,显著提升了检测速度。YOLOv5模型具有多种尺寸(如n、s、m、l、x),适用于不同的应用场景。本文将以YOLOv5n模型为例,展示如何在C++中进行推理部署。

3. 环境搭建

3.1 安装依赖库

  • OpenCV:用于图像处理和显示,提供了丰富的图像处理函数和数据结构。
  • ONNX Runtime:ONNX模型的推理引擎,负责加载和执行ONNX模型。

3.2 准备YOLOv5模型

将训练好的YOLOv5模型转换为ONNX格式,并确保模型文件路径正确。在代码中,通过指定模型文件路径来加载模型。

4. 代码实现

4.1 图像预处理

图像预处理是目标检测的重要环节,包括图像缩放、归一化等操作。在代码中,resize_image函数实现了图像的缩放和填充,保持图像的长宽比,使其适应模型的输入尺寸(640x640)。

Mat resize_image(Mat &srcimg, int &newh, int &neww, int &srch, int &srcw, int &top, int &left) {
   
    int inpHeight = 640;
    int inpWidth = 640;
    bool keep_ratio = true;
    Mat dstimg;
    if (keep_ratio && srch!= srcw) {
   
        // 计算缩放比例
        float hw_scale = (float)srch / srcw;
        if (hw_scale > 1) {
   
            newh = inpHeight;
            neww = int(inpWidth / hw_scale);
            resize(srcimg, dstimg, Size(neww, newh), INTER_AREA);
            left = int((inpWidth - neww) * 0.5);
            // 填充图像
            copyMakeBorder(dstimg, dstimg, 0, 0, left, inpWidth - neww - left, BORDER_CONSTANT, 114);
        } else {
   
            newh = (int)inpHeight * hw_scale;
            neww = inpWidth;
            resize(srcimg, dstimg, Size(neww, newh), INTER_AREA);
            top = (int)(inpHeight - newh) * 0.5;
            copyMakeBorder(dstimg, dstimg, top, inpHeight - newh - top, 0, 0, BORDER_CONSTANT, 114);
        }
    } else {
   
        resize(srcimg, dstimg, Size(neww, newh), INTER_AREA);
    }
    return dstimg;
}

4.2 模型推理

模型推理部分是核心,通过ONNX Runtime加载模型,并将预处理后的图像数据输入模型,获取预测结果。在infer函数中,首先初始化ONNX Runtime环境和会话,然后将图像数据转换为模型输入张量,执行推理,最后对输出结果进行后处理,包括解析边界框、置信度和类别信息。

void infer(cv::Mat frame, const std::string& onnxfilepath, float confThreshold, float iouThreshold,
            std::vector<cv::Rect>& boxes, std::vector<float>& confs, std::vector<int>& classIds, std::vector<int>& indices) {
   
    // 初始化相关变量
    if(boxes.size()) boxes.clear();
    if(confs.size()) confs.clear();
    if (classIds.size()) classIds.clear();
    if (indices.size()) indices.clear();
    static bool is_initialized = false;
    static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "yolov5s-5.0");
    static Ort::SessionOptions session_options;
    static std::vector<const char*> input_node_names = {
    "images" };
    static std::vector<const char*> output_node_names = {
    "output0" };
    static size_t input_tensor_size = 3 * 640 * 640;
    static std::vector<float> input_tensor_values(input_tensor_size);

    // 初始化模型路径
#ifdef _WIN32
    static wchar_t *model_path = 

你可能感兴趣的:(c++,YOLO,目标检测)