openCV调用YOLOv5 ONNX模型
opencv
OpenCV: 开源计算机视觉库
项目地址：https://gitcode.com/gh_mirrors/opencv31/opencv
免费下载资源
·
https://blog.csdn.net/qq_35135771/article/details/116592203
感谢大佬
导出U版YOLOv5的ONNX模型，openCV直接读取会失败，原因在于Pytorch2ONNX不支持对slice对象赋值，需要更改下列两个地方代码：
yolov5/models/common.py
-
class Focus(nn.Module):
-
# Focus wh information into c-space
-
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
-
super(Focus, self).__init__()
-
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
-
# self.contract = Contract(gain=2)
-
-
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
-
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
-
# return self.conv(self.contract(x))
-
-
更改为：
-
-
-
class Focus(nn.Module):
-
# Focus wh information into c-space
-
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
-
super(Focus, self).__init__()
-
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
-
# self.contract = Contract(gain=2)
-
-
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
-
# return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
-
N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
-
s = 2
-
x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)
-
x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
-
y = x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)
-
return self.conv(y)
具体原因可看博客2021.03.11更新 c++下使用opencv部署yolov5模型（一）_爱晚乏客游的博客-CSDN博客和Pytorch转ONNX-实战篇2（实战踩坑总结） - 知乎，这里不做过多解释
yolov5/models/yolo.py
-
class
Detect(nn.Module):
-
stride =
None
# strides computed during build
-
export =
False
# onnx export
-
-
def
__init__(
self, nc=80, anchors=(), ch=()):
# detection layer
-
super(Detect, self).__init__()
-
self.nc = nc
# number of classes
-
self.no = nc +
5
# number of outputs per anchor
-
self.nl =
len(anchors)
# number of detection layers
-
self.na =
len(anchors[
0]) //
2
# number of anchors
-
self.grid = [torch.zeros(
1)] * self.nl
# init grid
-
a = torch.tensor(anchors).
float().view(self.nl, -
1,
2)
-
self.register_buffer(
'anchors', a)
# shape(nl,na,2)
-
self.register_buffer(
'anchor_grid', a.clone().view(self.nl,
1, -
1,
1,
1,
2))
# shape(nl,1,na,1,1,2)
-
self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na,
1)
for x
in ch)
# output conv
-
-
def
forward(
self, x):
-
# x = x.copy() # for profiling
-
z = []
# inference output
-
self.training |= self.export
-
for i
in
range(self.nl):
-
x[i] = self.m[i](x[i])
# conv
-
bs, _, ny, nx = x[i].shape
# x(bs,255,20,20) to x(bs,3,20,20,85)
-
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(
0,
1,
3,
4,
2).contiguous()
-
x[i] = x[i].view(bs * self.na * ny * nx, self.no).contiguous()
-
return torch.cat(x)
-
# if not self.training: # inference
-
# if self.grid[i].shape[2:4] != x[i].shape[2:4]:
-
# self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
-
#
-
# y = x[i].sigmoid()
-
# y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
-
# y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
-
# z.append(y.view(bs, -1, self.no))
-
#
-
# return x if self.training else (torch.cat(z, 1), x)
更改完成之后，使用netron查看导出的onnx模型，可发现上述操作把YOLO层的三个输出合并为一个输出，格式类似yolov4的格式。
-
python
-
import netron
-
netron.start(
"best.onnx")
接下来就是使用openCV调用导出的onnx，本文使用的openCV版本为4.5.1，使用4.4版本的可能会出现错误。代码参考了博客2021.09.02更新说明 c++下使用opencv部署yolov5模型 （三）_爱晚乏客游的博客-CSDN博客，我也是依葫芦画瓢
yolo5.h
-
#pragma once
-
#include <fstream>
-
#include <sstream>
-
#include <iostream>
-
#include <opencv2/dnn.hpp>
-
#include <opencv2/imgproc.hpp>
-
#include <opencv2/highgui.hpp>
-
#include<time.h>
-
using
namespace cv;
-
using
namespace dnn;
-
using
namespace std;
-
struct
Output {
-
int id;
//结果类别id
-
float confidence;
//结果置信度
-
cv::Rect box;
//矩形框
-
};
-
-
class
YOLO
-
{
-
public:
-
YOLO() {};
-
bool readModel(Net &net, string &netPath, bool isCuda = false);
-
bool Detect(cv::Mat &SrcImg, cv::dnn::Net &net, std::vector<Output> &output);
-
void drawPred(cv::Mat &img, std::vector<Output> result, std::vector<cv::Scalar> color);
-
private:
-
//计算归一化函数
-
float Sigmoid(float x) {
-
return
static_cast<
float>(
1.f / (
1.f +
exp(-x)));
-
}
-
const
float netAnchors[
3][
6] = { {
17.0,
22.0,
23.0,
32.0,
31.0,
44.0 },{
40.0,
60.0,
51.0,
85.0,
68.0,
117.0 },{
108.0,
93.0,
94.0,
168.0,
170.0,
331.0 } };
-
const
float netStride[
3] = {
8.0,
16.0,
32.0 };
-
const
int netWidth =
608;
-
const
int netHeight =
608;
-
float nmsThreshold =
0.2;
-
float boxThreshold =
0.5;
-
float classThreshold =
0.5;
-
std::vector<std::string> className = {
"head",
"tail" };
-
};
yolo5.cpp
-
#include "yolo5.h"
-
#include<iostream>
-
#include<time.h>
-
using
namespace cv;
-
using
namespace std;
-
bool YOLO::readModel(Net & net, string & netPath, bool isCuda)
-
{
-
try {
-
net =
readNetFromONNX(netPath);
-
}
-
catch (
const std::exception&) {
-
return
false;
-
}
-
//cuda
-
if (isCuda) {
-
net.
setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
-
net.
setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
-
}
-
//cpu
-
else {
-
net.
setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
-
net.
setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
-
}
-
return
true;
-
}
-
-
bool YOLO::Detect(cv::Mat & SrcImg, cv::dnn::Net & net, std::vector<Output>& output)
-
{
-
Mat blob;
-
int col = SrcImg.cols, row = SrcImg.rows;
-
int maxLen = std::
max(col,row);
-
Mat netInputImg = SrcImg.
clone();
-
if (maxLen >
1.2*col || maxLen >
1.2*row) {
//纠正之前边框检测不准确的问题(2022-06-31)
-
Mat resizeImg = Mat::
zeros(maxLen,maxLen,CV_8UC3);
-
SrcImg.
copyTo(
resizeImg(
Rect(
0,
0,col,row)));
-
netInputImg = resizeImg.
clone();
-
}
-
blobFromImage(SrcImg, blob,
1 /
255.0, cv::
Size(netWidth, netHeight),
Scalar(
0,
0,
0),
true,
false);
-
net.
setInput(blob);
-
vector<Mat> netOutputImg;
-
net.forward(netOutputImg, net.
getUnconnectedOutLayersNames());
-
vector<
int> classIds;
//结果id数组
-
vector<
float> confidences;
//结果每个id对应置信度数组
-
vector<Rect> boxes;
//每个id矩形框
-
float ratio_h = (
float)netInputImg.rows / netHeight;
//纠正
-
float ratio_w = (
float)netInputImg.cols / netWidth;
//纠正
-
int net_width = className.
size() +
5;
//输出的网络宽度是类别数+5
-
float* pdata = (
float*)netOutputImg[
0].data;
-
for (
int stride =
0; stride <
3; stride++) {
//stride
-
int grid_x = (
int)(netWidth / netStride[stride]);
-
int grid_y = (
int)(netHeight / netStride[stride]);
-
for (
int anchor =
0; anchor <
3; anchor++) {
//anchors
-
const
float anchor_w = netAnchors[stride][anchor *
2];
-
const
float anchor_h = netAnchors[stride][anchor *
2 +
1];
-
for (
int i =
0; i < grid_x; i++) {
-
for (
int j =
0; j < grid_y; j++) {
-
float box_score =
Sigmoid(pdata[
4]);
//获取每一行的box框中含有某个物体的概率
-
if (box_score > boxThreshold) {
-
cv::Mat scores(1, className.size(), CV_32FC1, pdata + 5);
-
Point classIdPoint;
-
double max_class_socre;
-
minMaxLoc(scores,
0, &max_class_socre,
0, &classIdPoint);
-
max_class_socre =
Sigmoid((
float)max_class_socre);
-
if (max_class_socre > boxThreshold) {
-
//rect [x,y,w,h]
-
float x = (
Sigmoid(pdata[
0]) *
2.f -
0.5f + j) * netStride[stride];
//x
-
float y = (
Sigmoid(pdata[
1]) *
2.f -
0.5f + i) * netStride[stride];
//y
-
float w =
powf(
Sigmoid(pdata[
2]) *
2.f,
2.f) * anchor_w;
//w
-
float h =
powf(
Sigmoid(pdata[
3]) *
2.f,
2.f) * anchor_h;
//h
-
int left = (x -
0.5*w)*ratio_w;
-
int top = (y -
0.5*h)*ratio_h;
-
classIds.
push_back(classIdPoint.x);
-
confidences.
push_back(max_class_socre);
-
boxes.
push_back(
Rect(left, top,
int(w*ratio_w),
int(h*ratio_h)));
-
}
-
}
-
pdata += net_width;
//指针移到下一行
-
}
-
}
-
}
-
}
-
vector<
int> nms_result;
-
NMSBoxes(boxes, confidences, classThreshold, nmsThreshold, nms_result);
-
for (
int i =
0; i < nms_result.
size(); i++) {
-
int idx = nms_result[i];
-
Output result;
-
result.id = classIds[idx];
-
result.confidence = confidences[idx];
-
result.box = boxes[idx];
-
output.
push_back(result);
-
}
-
if (output.
size())
-
return
true;
-
else
-
return
false;
-
}
-
-
void YOLO::drawPred(Mat &img, vector<Output> result, vector<Scalar> color) {
-
for (
int i =
0; i < result.
size(); i++)
-
{
-
int left, top;
-
left = result[i].box.x;
-
top = result[i].box.y;
-
int color_num = i;
-
rectangle(img, result[i].box, color[result[i].id],
2,
8);
-
string label = className[result[i].id] +
":" +
to_string(result[i].confidence);
-
int baseLine;
-
Size labelSize =
getTextSize(label, FONT_HERSHEY_SIMPLEX,
0.5,
1, &baseLine);
-
top =
max(top, labelSize.height);
-
putText(img, label,
Point(left, top), FONT_HERSHEY_SIMPLEX,
1, color[result[i].id],
1);
-
}
-
cv::
resize(img, img, cv::
Size(
1400,
800));
-
imshow(
"test", img);
-
waitKey();
-
}
main.cpp
-
#include <opencv2/opencv.hpp>
-
#include <opencv2/dnn.hpp>
-
#include <highgui/highgui_c.h>
-
#include <iostream>
-
#include "yolo5.h"
-
using
namespace cv;
-
using
namespace cv::dnn;
-
using
namespace std;
-
int main()
-
{
-
string img_path =
"./test.jpg";
-
string model_path =
"./best.onnx";
-
YOLO test;
-
Net net;
-
if (test.
readModel(net, model_path,
false)) {
-
cout <<
"read net ok!" << endl;
-
}
-
else {
-
return
-1;
-
}
-
//生成随机颜色
-
vector<Scalar> color;
-
srand(
time(
0));
-
for (
int i =
0; i <
80; i++) {
-
int b =
rand() %
256;
-
int g =
rand() %
256;
-
int r =
rand() %
256;
-
color.
push_back(
Scalar(b, g, r));
-
}
-
vector<Output> result;
-
Mat img =
imread(img_path);
-
if (test.
Detect(img, net, result)) {
-
test.
drawPred(img, result, color);
-
}
-
else {
-
cout <<
"Detect Failed!" << endl;
-
}
-
system(
"pause");
-
return
0;
-
}
我做了一个夜间车辆检测，下面贴出运行结果。感觉检测结果边框有点变形，回头再看看哪出错了。模型文件可在我的博客资源里下载yolov5夜间车辆检测模型-深度学习文档类资源-CSDN文库
GitHub 加速计划 / opencv31 / opencv
77.34 K
55.71 K
下载
OpenCV: 开源计算机视觉库
最近提交(Master分支：1 个月前 )
e1fec156
features2d: fixed out of bounds access in SIFT 15 天前
63087396 - 15 天前
旨在为数千万中国开发者提供一个无缝且高效的云端环境，以支持学习、使用和贡献开源项目。
更多推荐
0
0- 0
已为社区贡献9条内容
所有评论(0)