使用B导的yolov7代码部署，代码地址：https://github.com/bubbliiiing/yolov7-pytorch
模型的的训练看B导即可，up主地址：Bubbliiiing的博客_CSDN博客-神经网络学习小记录,睿智的目标检测,有趣的数据结构算法领域博主
模型训练完成之后，在predict.py中设置mode = "export_onnx"即可生成。
注意，此处有个坑，B导的yolov7代码输出的onnx只有1*class.size*20*20这一层，需要在nets/yolo.py文件中修改一下。
修改之前：（在yolo.py的最下面）
#---------------------------------------------------#
# 第三个特征层
# y3=(batch_size, 36, 80, 80)
#---------------------------------------------------#
out2 = self.yolo_head_P3(P3)
#---------------------------------------------------#
# 第二个特征层
# y2=(batch_size, 36, 40, 40)
#---------------------------------------------------#
out1 = self.yolo_head_P4(P4)
#---------------------------------------------------#
# 第一个特征层
# y1=(batch_size, 36, 20, 20)
#---------------------------------------------------#
out0 = self.yolo_head_P5(P5)
return [out0, out1, out2]
修改之后：
#---------------------------------------------------#
# 第三个特征层
# y3=(batch_size, 36, 80, 80)
#---------------------------------------------------#
out2 = self.yolo_head_P3(P3)
bs, _, ny, nx = out2.shape # x(bs,255,20,20) to x(bs,3,20,20,85)
out2 = out2.view(bs, 3, 12, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
out2 = out2.view(bs * 3 * ny * nx, 12).contiguous()
#---------------------------------------------------#
# 第二个特征层
# y2=(batch_size, 36, 40, 40)
#---------------------------------------------------#
out1 = self.yolo_head_P4(P4)
bs, _, ny, nx = out1.shape
out1 = out1.view(bs, 3, 12, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
out1 = out1.view(bs * 3 * ny * nx, 12).contiguous()
#---------------------------------------------------#
# 第一个特征层
# y1=(batch_size, 36, 20, 20)
#---------------------------------------------------#
out0 = self.yolo_head_P5(P5)
bs, _, ny, nx = out0.shape
out0 = out0.view(bs, 3, 12, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
out0 = out0.view(bs * 3 * ny * nx, 12).contiguous()
#return [out0, out1, out2]
return torch.cat((out2,out1,out0))
这样我们可以看到输出的shape已经变成了25200*12了！
解释下这个数据：网络原本的输出是1*36*80*80，1*36*40*40，1*36*20*20，36是因为我的模型的类别数为7，36=（5+7）*3，5为四个位置信息加置信度，3为anchor数，经过上述代码的操作就把所有的输出拼接起来了，结果为25200*12，一共有25200个预测结果与每个结果对应12个信息。
之后我们就可以利用生成的onnx在vs studio中进行部署啦。
main.cpp：
#include "yolo.h"
#include
#include
#include
using namespace std;
using namespace cv;
using namespace dnn;
int main()
{
string img_path = "3.jpg";
string model_path = "models.onnx";
Yolo test;
Net net;
//加载onnx模型
if (test.readModel(net, model_path, true)) {
cout << "read net ok!" << endl;
}
else {
return -1;
}
//生成随机颜色
vector 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
yolo.h：
#pragma once
#include
#include
#include
struct Output {
int id;//结果类别id
float confidence;//结果置信度
cv::Rect box;//矩形框
};
class Yolo
{
public:
Yolo() {}
~Yolo(){}
bool readModel(cv::dnn::Net& net, std::string& netPath, bool isCuda);
bool Detect(cv::Mat& SrcImg, cv::dnn::Net& net, std::vector
yolo.cpp:
#include "Yolo.h"
using namespace std;
using namespace cv;
using namespace dnn;
bool Yolo::readModel(Net& net, string& netPath, bool isCuda = false) {
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(Mat& SrcImg, Net& net, vector
预测结果：
大功告成啦，不得不说yolov7的效果相当的好，也感谢B导大大啦
