赞
踩
导出U版YOLOv5-V5.0版本(V6.0以上版本已经不用Focus层了,直接跳过这步)的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文库
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。