在win10+opencv4.4.0+opencv_contrib4.4.0+cuda10.2环境下部署yolov5模型
编译过程参考了这一篇博客
https://blog.csdn.net/hai_fellow_Z/article/details/116002203
主要有两个代码yolo.h和main_yolo.cpp
主程序为main_yolo.cpp文件
以下代码main函数可以分别实现图片、视频、摄像头的调用
#include "yolo.h"
#include <opencv2/opencv.hpp>
using namespace cv;
YOLO::YOLO(Net_config config)
{
cout << "Net use " << config.netname << endl;
this->confThreshold = config.confThreshold;
this->nmsThreshold = config.nmsThreshold;
this->objThreshold = config.objThreshold;
strcpy_s(this->netname, config.netname.c_str());
ifstream ifs(this->classesFile.c_str());
string line;
while (getline(ifs, line)) this->classes.push_back(line);
string modelFile = this->netname;
modelFile += ".onnx";
modelFile = "F:\\code\\a14-yolov5-dnn-opencv\\x64\\Release\\yolov5s.onnx";
this->net = readNet(modelFile);
net.setPreferableTarget(DNN_TARGET_CUDA);
net.setPreferableBackend(DNN_BACKEND_CUDA);
}
void YOLO::drawPred(int classId, float conf, int left, int top, int right, int bottom, Mat & frame)
{
rectangle(frame, Point(left, top), Point(right, bottom), Scalar(0, 0, 255), 3);
string label = format("%.2f", conf);
label = this->classes[classId] + ":" + label;
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = max(top, labelSize.height);
putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
}
void YOLO::sigmoid(Mat * out, int length)
{
float* pdata = (float*)(out->data);
int i = 0;
for (i = 0; i < length; i++)
{
pdata[i] = 1.0 / (1 + expf(-pdata[i]));
}
}
void YOLO::detect(Mat & frame)
{
Mat blob;
blobFromImage(frame, blob, 1 / 255.0, Size(this->inpWidth, this->inpHeight), Scalar(0, 0, 0), true, false);
this->net.setInput(blob);
vector<Mat> outs;
std::vector<String> res = this->net.getUnconnectedOutLayersNames();
this->net.forward(outs, res);
vector<int> classIds;
vector<float> confidences;
vector<Rect> boxes;
float ratioh = (float)frame.rows / this->inpHeight, ratiow = (float)frame.cols / this->inpWidth;
int n = 0, q = 0, i = 0, j = 0, nout = this->classes.size() + 5, c = 0;
for (n = 0; n < 3; n++)
{
int num_grid_x = (int)(this->inpWidth / this->stride[n]);
int num_grid_y = (int)(this->inpHeight / this->stride[n]);
int area = num_grid_x * num_grid_y;
this->sigmoid(&outs[n], 3 * nout * area);
for (q = 0; q < 3; q++)
{
const float anchor_w = this->anchors[n][q * 2];
const float anchor_h = this->anchors[n][q * 2 + 1];
float* pdata = (float*)outs[n].data + q * nout * area;
for (i = 0; i < num_grid_y; i++)
{
for (j = 0; j < num_grid_x; j++)
{
float box_score = pdata[4 * area + i * num_grid_x + j];
if (box_score > this->objThreshold)
{
float max_class_socre = 0, class_socre = 0;
int max_class_id = 0;
for (c = 0; c < this->classes.size(); c++)
{
class_socre = pdata[(c + 5) * area + i * num_grid_x + j];
if (class_socre > max_class_socre)
{
max_class_socre = class_socre;
max_class_id = c;
}
}
if (max_class_socre > this->confThreshold)
{
float cx = (pdata[i * num_grid_x + j] * 2.f - 0.5f + j) * this->stride[n];
float cy = (pdata[area + i * num_grid_x + j] * 2.f - 0.5f + i) * this->stride[n];
float w = powf(pdata[2 * area + i * num_grid_x + j] * 2.f, 2.f) * anchor_w;
float h = powf(pdata[3 * area + i * num_grid_x + j] * 2.f, 2.f) * anchor_h;
int left = (cx - 0.5 * w) * ratiow;
int top = (cy - 0.5 * h) * ratioh;
classIds.push_back(max_class_id);
confidences.push_back(max_class_socre);
boxes.push_back(Rect(left, top, (int)(w * ratiow), (int)(h * ratioh)));
}
}
}
}
}
}
vector<int> indices;
NMSBoxes(boxes, confidences, this->confThreshold, this->nmsThreshold, indices);
for (size_t i = 0; i < indices.size(); ++i)
{
int idx = indices[i];
Rect box = boxes[idx];
this->drawPred(classIds[idx], confidences[idx], box.x, box.y,
box.x + box.width, box.y + box.height, frame);
}
}
int main()
{
YOLO yolo_model(yolo_nets[0]);
cv::VideoCapture cap;
cap.open("F:/code/yolov5_c++_dnn/yolov5/yolov5/ren.mp4");
if (!cap.isOpened()) {
std::cout << "打不开"<< std::endl;
return 0;
}
Mat frame;
while (1)
{
cap >> frame;
if (frame.empty())
break;
yolo_model.detect(frame);
imshow("video", frame);
waitKey(1);
}
cap.release();
return 0;
}
头文件为yolo.h
#include <fstream>
#include <sstream>
#include <iostream>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
using namespace cv;
using namespace dnn;
using namespace std;
struct Net_config
{
float confThreshold;
float nmsThreshold;
float objThreshold;
string netname;
};
class YOLO
{
public:
YOLO(Net_config config);
void detect(Mat& frame);
private:
const float anchors[3][6] = {{10.0, 13.0, 16.0, 30.0, 33.0, 23.0}, {30.0, 61.0, 62.0, 45.0, 59.0, 119.0},{116.0, 90.0, 156.0, 198.0, 373.0, 326.0}};
const float stride[3] = { 8.0, 16.0, 32.0 };
const string classesFile = "F:\\code\\a14-yolov5-dnn-opencv\\coco.names";
const int inpWidth = 640;
const int inpHeight = 640;
float confThreshold;
float nmsThreshold;
float objThreshold;
char netname[20];
vector<string> classes;
Net net;
void drawPred(int classId, float conf, int left, int top, int right, int bottom, Mat& frame);
void sigmoid(Mat* out, int length);
};
static inline float sigmoid_x(float x)
{
return static_cast<float>(1.f / (1.f + exp(-x)));
}
Net_config yolo_nets[4] = {
{0.5, 0.5, 0.5, "yolov5s"},
{0.5, 0.5, 0.5, "yolov5m"},
{0.5, 0.5, 0.5, "yolov5l"},
{0.5, 0.5, 0.5, "yolov5x"}
};
图片测试效果(注:视频资源和摄像头的调用可以自己调试)
|