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 =