c++ onnx之yolov5检测
·
yolov5和resnet比稍微麻烦了一点，主要就是多了nms部分，还有坐标点映射回原图的yolov5_scale_coords函数。流程大致分为五部分：1）图像等比例放缩，2）图像预处理，3）onnx推理，4）nms后处理，5）坐标点映射回原图
等比例放缩
还是和resnet一样的 letterbox 函数，就不重复了。
图像预处理
还是和resnet一样的，就不重复了，就cv::dnn::blobFromImage一句话。
onnx推理
和之前差不多，改下输入输出的维度、名称，然后resnet的一维输出变成yolo的二维输出，这个在 c++ 使用onnx推理 中写过了，也不重复了。
nms后处理
这个我懒得自己写，还好opencv里面有cv::dnn::NMSBoxes函数，直接用就好。
正常来说onnx推理后返回的预测值是252000 x (n+5) ，其中252000是框的总数一般是不变的， n是等于你设置的类别数，向量含义是（center_x, center_y, width, height, conf_框, conf_类别1，conf_类别2，…,，conf_类别n）。
这个函数主要就是把preds转成opencv_nms函数需要的输入 cv::Rect(left, top, width, height) 框坐标 和 float 置信度（类别置信度*框置信度）
vector<Detection> yolov5_nms(Mat preds, float conf_thres = 0.25, float iou_thres = 0.45)
{
vector<cv::Rect> boxes;
vector<float> confs;
vector<int> classIds;
for (int i = 0; i < preds.rows; i++)
{
float clsConf = preds.at<float>(i, 4);
if (clsConf > conf_thres)
{
// (cx,cy,w,h) to (left,top,w,h)
float centerX = preds.at<float>(i, 0);
float centerY = preds.at<float>(i, 1);
float width = preds.at<float>(i, 2);
float height = preds.at<float>(i, 3);
float left = centerX - width / 2;
float top = centerY - height / 2;
// 因为我这里只检测人，就直接这样来了，正常如果有80个类别，objConf表示最高的置信度，classId表示最高置信度的id
float objConf = preds.at<float>(i, 5);;
int classId = 0;
float confidence = clsConf * objConf;
boxes.push_back(cv::Rect(left, top, width, height));
confs.push_back(confidence);
classIds.push_back(classId);
}
}
vector<int> nms_result;
// 这里输入的boxes(float)被强制转为(int)了，可能会有点误差
cv::dnn::NMSBoxes(boxes, confs, conf_thres, iou_thres, nms_result);
cout << "amount of NMS indices: " << nms_result.size() << std::endl;
vector<Detection> output;
for (int i = 0; i < nms_result.size(); i++) {
int idx = nms_result[i];
Detection result;
result.class_id = classIds[idx];
result.confidence = confs[idx];
result.box = boxes[idx];
output.push_back(result);
}
return output;
}
坐标点映射
主要就是我要把检测的框放缩回原图，对原图进行裁减。结构体看着有点难受，可以直接输入输出改成boxes，我是因为之后可视化要标上类别信息所以写了个结构体。
vector<Detection> yolov5_scale_coords(Mat ori_img, vector<Detection> detections, Mat letter_img)
{
// boxes(ltwh)(left, top, width, height)
float ratio;
float scale_row = (float)(letter_img.rows) / (float)(ori_img.rows);
float scale_col= (float)(letter_img.cols) / (float)(ori_img.cols);
ratio = min(scale_row, scale_col);
float pad_row = ((float)(letter_img.rows) - (float)(ori_img.rows) * ratio) / 2.0;
float pad_col = ((float)(letter_img.cols) - (float)(ori_img.cols) * ratio) / 2.0;
// 这里的点应该要用float的，用int明显有误差
vector<Detection> new_detections;
int nums = detections.size();
for (int i = 0; i < nums; ++i)
{
auto detection = detections[i];
auto box = detection.box;
float left = (box.x - pad_col) / ratio;
float top = (box.y - pad_row) / ratio;
float width = box.width / ratio;
float height = box.height / ratio;
// 这种情况可能出现越界，不处理的话之后截图会出问题
if (left < 0) {
width = width + left;
left = -left;
}
if (top < 0) {
height = height + top;
top = -top;
}
height = min(height, ori_img.rows - top);
width = min(width, ori_img.cols - left);
Detection new_detection;
new_detection.class_id = detection.class_id;
new_detection.confidence = detection.confidence;
new_detection.box = cv::Rect(left, top, width, height);
new_detections.push_back(new_detection);
}
return new_detections;
}
汇总
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc_c.h>
#include <opencv2/dnn.hpp>
#include <iostream>
#include <onnxruntime_cxx_api.h>
#include <assert.h>
#include <vector>
#include <fstream>
using namespace cv; //当定义这一行后，cv::imread可以直接写成imread
using namespace std;
using namespace Ort;
using namespace cv::dnn;
#define IMG_LEN 640
#define pi 3.1415926
struct Detection
{
int class_id;
float confidence;
cv::Rect box;
};
Mat letterbox(Mat src)
{
//以下为带边框图像生成
int in_w = src.cols;
int in_h = src.rows;
int tar_w = IMG_LEN;
int tar_h = IMG_LEN;
//哪个缩放比例小选用哪个
float r = min(float(tar_h) / in_h, float(tar_w) / in_w);
int inside_w = round(in_w * r);
int inside_h = round(in_h * r);
int padd_w = tar_w - inside_w;
int padd_h = tar_h - inside_h;
//内层图像resize
Mat resize_img;
resize(src, resize_img, Size(inside_w, inside_h));
//cvtColor(resize_img, resize_img, COLOR_BGR2RGB);
padd_w = padd_w / 2;
padd_h = padd_h / 2;
//外层边框填充灰色
int top = int(round(padd_h - 0.1));
int bottom = int(round(padd_h + 0.1));
int left = int(round(padd_w - 0.1));
int right = int(round(padd_w + 0.1));
copyMakeBorder(resize_img, resize_img, top, bottom, left, right, BORDER_CONSTANT, Scalar(114, 114, 114));
//cout << resize_img.size() << endl;
//imshow("pad", resize_img);
//waitKey(10);
return resize_img;
}
Mat yolov5_onnx_model(Mat blob)
{
//#ifdef _WIN32
// wstring model_path = charToWstring(modelPath.c_str());
//#else
// string model_path = modelPath.c_str();
//#endif
#ifdef _WIN32
const wchar_t* model_path = L"E://c++//mmpose//yolov5s.onnx";
#else
const char* model_path = "E://c++//mmpose//yolov5s.onnx";
#endif
//environment （设置为VERBOSE（ORT_LOGGING_LEVEL_VERBOSE）时，方便控制台输出时看到是使用了cpu还是gpu执行）
Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "OnnxModel");
Ort::SessionOptions session_options;
// 使用1个线程执行op,若想提升速度，增加线程数
session_options.SetIntraOpNumThreads(1);
//CUDA加速开启(由于onnxruntime的版本太高，无cuda_provider_factory.h的头文件，加速可以使用onnxruntime V1.8的版本)
//OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
// ORT_ENABLE_ALL: 启用所有可能的优化
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
//load model and creat session
//printf("Using Onnxruntime C++ API\n");
Ort::Session session(env, model_path, session_options);
// print model input layer (node names, types, shape etc.)
Ort::AllocatorWithDefaultOptions allocator;
//model info
// 获得模型又多少个输入和输出，一般是指对应网络层的数目
// 一般输入只有图像的话input_nodes为1
size_t num_input_nodes = session.GetInputCount();
// 如果是多输出网络，就会是对应输出的数目
size_t num_output_nodes = session.GetOutputCount();
//printf("Number of inputs = %zu\n", num_input_nodes);
//printf("Number of output = %zu\n", num_output_nodes);
//获取输入name
const char* input_name = session.GetInputName(0, allocator);
//std::cout << "input_name:" << input_name << std::endl;
//获取输出name
const char* output_name = session.GetOutputName(0, allocator);
//std::cout << "output_name: " << output_name << std::endl;
// 自动获取维度数量
auto input_dims = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
auto output_dims = session.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
//std::cout << "input_dims:" << input_dims[0] << std::endl;
//std::cout << "output_dims:" << output_dims[0] << std::endl;
std::vector<const char*> input_names{ input_name };
std::vector<const char*> output_names = { output_name };
std::vector<const char*> input_node_names = { "images" }; //自己打开onnx模型看下,名称不一样会报错
std::vector<const char*> output_node_names = { "output" };
clock_t startTime, endTime;
//创建输入tensor
auto memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
std::vector<Ort::Value> input_tensors;
input_tensors.emplace_back(Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_dims.data(), input_dims.size()));
/*cout << int(input_dims.size()) << endl;*/
startTime = clock();
//推理(score model & input tensor, get back output tensor)
auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), input_tensors.data(), input_names.size(), output_node_names.data(), output_node_names.size());
endTime = clock();
assert(output_tensors.size() == 1 && output_tensors.front().IsTensor());
//除了第一个节点外，其他参数与原网络对应不上程序就会无法执行
//第二个参数代表输入节点的名称集合
//第四个参数1代表输入层的数目
//第五个参数代表输出节点的名称集合
//最后一个参数代表输出节点的数目
//获取输出(Get pointer to output tensor float values)
Ort::Value& det_out = output_tensors.at(0);
// 得到最可能分类输出
Mat newarr = Mat_<float>(25200, 7); //定义一个1*1000的矩阵
for (int i = 0; i < newarr.rows; i++)
{
vector<float> out;
for (int j = 0; j < newarr.cols; j++) //矩阵列数循环
{
newarr.at<float>(i, j) = det_out.At<float>({ 0, i, j });
}
}
//for (int i = 0; i < newarr.rows; i++)
//{
// for (int j = 0; j < newarr.cols; j++) //矩阵列数循环
// {
// cout<< out_list[i][j]<<", ";
// }
// cout << endl;
//}
//cout << newarr.size() << endl;
return newarr;
}
vector<Detection> yolov5_nms(Mat preds, float conf_thres = 0.25, float iou_thres = 0.45)
{
vector<cv::Rect> boxes;
vector<float> confs;
vector<int> classIds;
for (int i = 0; i < preds.rows; i++)
{
float clsConf = preds.at<float>(i, 4);
if (clsConf > conf_thres)
{
float centerX = preds.at<float>(i, 0);
float centerY = preds.at<float>(i, 1);
float width = preds.at<float>(i, 2);
float height = preds.at<float>(i, 3);
float left = centerX - width / 2;
float top = centerY - height / 2;
// 因为我这里只检测人，就直接这样来了，正常如果有80个类别，objConf表示最高的置信度，classId表示最高置信度的id
float objConf = preds.at<float>(i, 5);;
int classId = 0;
float confidence = clsConf * objConf;
boxes.push_back(cv::Rect(left, top, width, height));
confs.push_back(confidence);
classIds.push_back(classId);
}
}
vector<int> nms_result;
// 这里输入的boxes(float)被强制转为(int)了，可能会有点误差
cv::dnn::NMSBoxes(boxes, confs, conf_thres, iou_thres, nms_result);
cout << "amount of NMS indices: " << nms_result.size() << std::endl;
vector<Detection> output;
for (int i = 0; i < nms_result.size(); i++) {
int idx = nms_result[i];
Detection result;
result.class_id = classIds[idx];
result.confidence = confs[idx];
result.box = boxes[idx];
output.push_back(result);
}
return output;
}
void yolov5_imshow(Mat frame, vector<Detection> output, vector<string> class_list)
{
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) };
int nums = output.size();
for (int i = 0; i < nums; ++i)
{
auto detection = output[i];
auto box = detection.box;
auto classId = detection.class_id;
const auto color = colors[classId % colors.size()];
cv::rectangle(frame, box, color, 3);
cv::rectangle(frame, cv::Point(box.x, box.y - 20), cv::Point(box.x + box.width, box.y), color, cv::FILLED);
cv::putText(frame, class_list[classId].c_str(), cv::Point(box.x, box.y - 5), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
}
cv::imshow("output", frame);
waitKey(0);
}
vector<Detection> yolov5_scale_coords(Mat ori_img, vector<Detection> detections, Mat letter_img)
{
// boxes(ltwh)(left, top, width, height)
float ratio;
float scale_row = (float)(letter_img.rows) / (float)(ori_img.rows);
float scale_col= (float)(letter_img.cols) / (float)(ori_img.cols);
ratio = min(scale_row, scale_col);
float pad_row = ((float)(letter_img.rows) - (float)(ori_img.rows) * ratio) / 2.0;
float pad_col = ((float)(letter_img.cols) - (float)(ori_img.cols) * ratio) / 2.0;
// 这里的点应该要用float的，用int明显有误差
vector<Detection> new_detections;
int nums = detections.size();
for (int i = 0; i < nums; ++i)
{
auto detection = detections[i];
auto box = detection.box;
float left = (box.x - pad_col) / ratio;
float top = (box.y - pad_row) / ratio;
float width = box.width / ratio;
float height = box.height / ratio;
// 这种情况可能出现越界，不处理的话之后截图会出问题
if (left < 0) {
width = width + left;
left = -left;
}
if (top < 0) {
height = height + top;
top = -top;
}
height = min(height, ori_img.rows - top);
width = min(width, ori_img.cols - left);
Detection new_detection;
new_detection.class_id = detection.class_id;
new_detection.confidence = detection.confidence;
new_detection.box = cv::Rect(left, top, width, height);
new_detections.push_back(new_detection);
}
return new_detections;
}
vector<Detection> yolov5_main(Mat img, vector<string> class_list)
{
//图片预处理
Mat det1 = letterbox(img);
Mat det2 = dnn::blobFromImage(det1, 1. / 255, Size(IMG_LEN, IMG_LEN), Scalar(0.485, 0.456, 0.406), true, false);
printf("Processing img ......\n");
// 预测框
Mat preds = yolov5_onnx_model(det2);
// NMS
float conf_thres = 0.25;
float iou_thres = 0.45;
vector<Detection> output = yolov5_nms(preds, conf_thres, iou_thres);
// 坐标映射回原图
vector<Detection> new_output = yolov5_scale_coords(img, output, det1);
//// 显示图片
//yolov5_imshow(det1, output, class_list); //放缩之后的640x640
//yolov5_imshow(img, new_output, class_list); // 映射回原图
return new_output;
}
int main()
{
string imgpath = "E://c++//mmpose//test_yolo.jpg";
Mat img = imread(imgpath);
// 检测类别，我是只有person和其他两类
vector<string> class_list;
class_list.push_back("person");
class_list.push_back("other");
vector<Detection> detections = yolov5_main(img, class_list);
yolov5_imshow(img, detections, class_list);
//system("pause");
}
腾讯云面向开发者汇聚海量精品云计算使用和开发经验，营造开放的云计算技术生态圈。
更多推荐
0
0- 0
已为社区贡献4条内容
所有评论(0)