yolov5 训练模型可以使用CPU和GPU两种方式,CPU相较于GPU很慢,但是要使用GPU需要配置CUDA+cuDNN的环境。训练模型需要搭建pytorch的深度学习库环境。
配置CUDA+cuDNN,以及pytorch环境:
配置的教程网上已有许多,这里不再赘述。下面提供网址:
CUDA网址:https://developer.nvidia.com/cuda-downloads
cuDNN网址:https://developer.nvidia.com/rdp/cudnn-archive
Pytorch网址:PyTorchhttps://pytorch.org/
cuDNN下载好后解压,将其中的三个文件夹复制到安装CUDA的文件夹下,覆盖掉其中的内容,然后在环境变量中添加路径即可(如果存在多个版本,则将需要使用的版本移动到前面去即可,如下),如何进入环境变量,后面的【第四章——下载YOLOv5所需要的配套库】有详细步骤。
如何查看电脑可以支持的CUDA版本:
同时按下键盘的Win和R键,输入cmd,打开【命令提示符】,输入nvidia-smi
如何查看自己安装的CUDA版本,或者多个CUDA版本中当前使用的版本:
在【命令提示符】中输入nvcc -V
如何查看自己下载的pytorch版本是CUDA版本还是CPU版本:
这里提供Anaconda中创建虚拟环境,然后安装pytorch的查看方法。首先打开【命令提示符】,输入【conda activate 虚拟环境名称】(中括号只是标识,不输入)激活虚拟环境,然后输入【conda list】,如何创建虚拟环境,在后面的【第四章——下载YOLOv5所需要的配套库】有详细步骤。
-
二、下载YOLOv5完整代码
注意:如果需要像本篇文章一样,使用C++调用导出的torchscript模型,请下载yolov5 v5.0的版本,其中的文件内容如下:
Yolov5 v5.0的网盘链接:
链接:https://pan.baidu.com/s/1uNPRZWWqpkPB3lyRBX4AdQ
提取码:6666
如果需要其他的版本:
https://github.com/ultralytics/yolov5/tags
https://github.com/ultralytics/yolov5/releases
Yolov5 v5.0的GitHup下载界面
方法一:可以复制HTTPS方法提供的网址,然后打开【命令提示符】(快捷键Win+R,输入cmd回车),用git下载(前提:已经安装好git,没有安装git,也可以直接在git官网:Git - Downloads下载,然后直接安装,git环境是直接配置好的,安装路径可自行修改)。
方法二:也可以把这个仓库克隆到码云,从码云上下载。这个方法百度一下有很多博客写道。百度一下:“github迁移码云”这类关键字即可。下载完后,解压到想要的路径。
方法三:直接下载源码压缩包,然后解压即可,但是下载很慢,没有梯子不推荐此方式。
-
三、下载官方模型(权重文件)
-
1、说明
yolov5需要四个不同的权重文件,分别是yolov5s.pt、yolov5m.pt、yolov5l.pt和yolov5x.pt。
-
- yolov5s.,pt是最小的模型,适合在低端设备上实现快速的推理。它包含在COCO数据集上训练的80个不同的类别。
- yolov5m.pt是中型模型,也包含COCO数据集上的80个类别,但比yolov5s.pt的性能更好。它适用于中端设备,如笔记本电脑和普通的桌面计算机.
- yolov5l.pt是一个大型模型,适用于高端设备。它比yolov5m.pt具有更高的准确性和性能,但需要更高的计算资源。对于高端计算机和服务器,这是一个理想的选择。
- yolov5x.pt是最大的模型,而且是性能最好的。它比其他模型更快,同时在精度方面也要好得多。它需要大量的计算资源,所以只适用于高端服务器和GPU。
-
2、下载
1.下面是我提供的四个pt权重文件的网盘链接:
链接:https://pan.baidu.com/s/1aIhGc2O8182XUdiWGaM-Qw
提取码:6666
如果以上步骤没有找到,可以直接进入Releases · ultralytics/yolov5 · GitHub下载权重文件,注意版本。
2.在yolov5的文件夹里新建weights文件夹,将下载的权重文件放进去,此处位置没有强制要求,后续更改文件内容即可。
-
四、下载YOLOv5所需要的配套库
前提:已经安装Anaconda并配置好环境!!
Anaconda环境配置:
修改Anaconda虚拟环境的保存地址:
1.在C盘的用户文件夹下找到 .condarc 文件,如下
2.记事本打开,并添加虚拟环境的路径,如下
创建虚拟环境:
① 在【命令提示符】创建虚拟环境(环境名:pytorch,可自行命名)
conda create -n pytorch python==3.7 (创建虚拟环境还可以采取后续办法)
**在yolov5中尽量用python3.7,也可以偏上版本**。
② 进入环境
conda activate pytorch 或 activate pytorch
③ 下载安装配套库
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt
(使用清华镜像源)
在pip install的时候,可能会出现read timeout的情况,你需要更换镜像源,或者多执行几次pip install(使用pip命令时记得先关掉梯子)
可在cmd创建、激活和下载,也可以在Anaconda Prompt创建、激活和下载,操作同cmd一样。创建虚拟环境也可以使用Anaconda Navigator,操作如下:
-
五、检测YOLOv5的环境是否配置成功
修改detect.py里面的weights(权重)的路径,修改为自己下载的预训练权重文件的路径(也可以用记事本打开修改),这里也可以不修改,在后面cmd运行时直接输入对应命令即可,但如果是用pycharm等应用运行,则需要进行修改)。
在yolov5目录下运行cmd(前面演示过如何打开某个目录下的cmd)。
然后输入:
conda activate pytorch(激活虚拟环境,前面创建过的)
然后输入(没有修改detect.py文件):
python detect.py --source 0 --weights=“weights/yolov5s.pt”
(如果设备有摄像头)
或者
python detect.py --source=“data/images/bus.jpg” --weights=“weights/yolov5s.pt”
(设备没有摄像头)
前面修改过detect.py文件,这里可直接输入:
python detect.py --source 0 (如果设备有摄像头)
或者
python detect.py --source=“data/images/bus.jpg” (设备没有摄像头)
有摄像头的操作会打开摄像头并实时探测物品,没有摄像头的操作将用项目自带的一张测试图片进行测试。运行结束后,会打印结果文件的位置。在这个位置中可以找到测试结果。
测试图像的运行结果:
如果出现运行问题AttributeError: 'Detect' object has no attribute 'm',则说明权重文件有问题,需要选择其他权重文件。可以先尝试将14.5M的yolov5s.pt文件替换掉原来27M的yolov5s.pt文件。
-
六、训练自己的数据集
(一)准备数据集
1、使用labelimg标注自己的数据集
下载命令:pip install labelimg
(最好是在自己的虚拟环境中安装,python版本最好不要太高,python3.7到python3.9就好)
下面是labelimg的使用指南:
使用labelimg进行标注时可能会出现闪退的情况,用命令行激活labelimg进行标注,可以查看错误原因:
遇到该问题可能是 python 版本过高不发生自动类型转换
解决办法:
方法一:修改源代码
根据报错可知是参数不匹配,只要找到源码的位置修改内容即可。
源码位置在 D:\workspace\Anaconda\envs\pytorch\Lib\site-packages\libs\canvas.py 的第530行
在以下三行改变 drawRect,drawLine 的参数为 int
同理,若遇到类似问题也可以修改源代码使其正常运行
方法二:替换Python版本为python3.7-python3.9范围内的版本
创建新的虚拟环境,其python版本为3.7后即可正常使用labelimg
2、在yolov5中加入自己的数据集
在yolov5的data文件夹中创建自己的数据文件夹my_data,里面创建两个文件夹images(存放图片)和labels(存放标签)。
images文件夹:
里面分test(测试集)、train(训练集)、val(验证集)三个文件夹分别存储图片,train文件夹中存放大多数图片,test和val存放少量图片即可,如果只是训练模型,可不准备test文件夹。
labels文件夹:
里面分test、train、val三个文件夹分别存储txt标签(存储图片中检测对象的上下左右四个点位信息),和images文件夹中的分布相对应,且labels文件夹的三个子文件夹中均存在classes.txt文件(存储检测的类别信息)。
(二)文件修改
1、修改数据集方面的yaml文件
打开data/coco.yaml进行修改,可另存为自己命名的文件后修改。
2、修改网络参数方面的yaml文件
打开models文件夹下的yaml文件【当然,你想用哪个模型就去修改对应的yaml文件】
3、修改train.py中的一些参数
train.py在根目录里,修改一些主要的参数(也可不修改,后续训练时输入命令修改,缺点是每次训练都需要输入很长的命令进行修改)
(三)开始训练
- 打开yolov5的目录下的【命令提示符】
- 激活虚拟环境:conda activate pytorch
- 直接输入python train.py,也可选择如下方式输入,修改对应参数
例1:python train.py --data huahen.yaml --cfg yolov5s.yaml --weights weights/yolov5s.pt --epochs 10
(huahen.yaml:自己另存为的数据集文件)
(yolov5s.yaml:训练网络文件,和权重文件相对应)
(yolov5s.pt:权重文件)
(epochs 10:训练轮数)
例2:python train.py --data huahen.yaml --epochs 10 --weights ./weights/yolov5s.pt --cfg yolov5s.yaml --batch-size 8:
训练完后可以在runs/train/exp里找到自己的训练结果(多次训练,结果在exp1、exp2……)
ERROR:页面太小
解决方法:
- 调整虚拟内存大小(主要方法)
- 修改batch-size的值,调小一些(如果虚拟内存不足,可尝试调小一些)
如何调节虚拟内存:
右键桌面,点击“显示设置”,进入系统设置界面,在左边点击“关于”,进入点击“高级系统设置”,后续操作见下图。
-
七、导出Python训练的yolov5模型
在训练好的yolov5 pt 模型可以通过export.py进行导出 onnx和torchscript两种格式的权重模型文件。
在官方的文档中 说明了可以导出的具体的类型。
导出流程
1、修改export.py文件
如需使用本篇文章后续的C++调用torchscript模型的方法,则需要严格按照如下方式修改。
第一处:
第二处:
第三处:
如需使用onnx,此处可修改其版本:
注意:我导出的时候存在torchscript导出成功,但是onnx导出失败的问题,由于我只需要使用torchscript,所以没有再进行修改。
2、导出模型
控制台:
- 打开【命令提示符】,进入自己创建的虚拟环境
- 输入命令:python export.py
Pycharm等可运行Python的软件:
直接运行export.py文件。
八、用C++调用导出的torchscript模型
(一)环境配置
1、下载并解压C++版本的Torch
官网:PyTorchhttps://pytorch.org/
Torch的C++版本有CUDA版本和CPU版本,区别在于后面在VS中配置环境时是否配置CUDA的路径,是否能够使用GPU。
如果需要早期的libtorch版本,可以参考这篇博客:
如果还是需要其他版本的libtorch,可以打开梯子,看看这个网址的内容:
https://github.com/pytorch/pytorch/issues/40961
解压后的LibTorch文件夹:
2、在Visual Studio2022上配置LibTorch(类似于配置OpenCV)
打开VS2022,在项目上右键打开属性,依次进行如下配置:
1、VC++目录下的包含目录
libtorch解压目录\include
libtorch解压目录\include\torch\csrc\api\include
CUDA版本的LibTorch需额外配置:
CUDA安装目录\v12.1\include
2、VC++目录下的库目录
CUDA版本和CPU版本的PyTorch均需配置:
libtorch解压目录\lib
CUDA版本的LibTorch需额外配置:
CUDA安装目录\v12.1\lib\X64
3、链接器下的输入
CUDA版本和CPU版本的LibTorch均需配置(二选其一):
libtorch解压目录\lib\*.lib(所有lib文件均加入项目,程序所占空间较大)
文件名.lib(只加入自己所需的lib文件)
我这里加入的lib文件如下(仅供参考)
CUDA版本的LibTorch需额外配置(二选其一):
CUDA安装目录\ v12.1\lib\X64\*.lib
文件名.lib
3、C++部署Libtorch出现问题、错误汇总
博客:https://blog.csdn.net/zzz_zzz12138/article/details/109138805
1、由于找不到xxx.dll,无法继续执行代码,重新安装程序可能会解决此问题
(1)由于找不到c10.dll(或其他libtorch/lib中的.dll动态库),无法继续执行代码
(2)由于找不到VCRUNTIME 140_1D.dll,无法继续执行代码
2、LINK : fatal error LNK1104: cannot open file 'torch-NOTFOUND.obj' (torch-NOTFOUND.obj无法找到)
3、error C2440: “初始化”: 无法从“`torch::jit::script::Module`”转换为“`std::shared_ptr`”
4、无法定位程序输入点cudnnSetCTCLossDescriptorEx于动态链接库xxx.dll上
5、有未经处理的异常:Microsoft C++异常:c10::Error,位于内存位置xxx处
6、引发异常:0xC0000005:读取位置0xFFFFFFFFFFFFFFFE时发生访问冲突
7、error :c2872 std 不明确的符号
补充:由于找不到c10.dll(或其他libtorch/lib中的.dll动态库),无法继续执行代码
把libtorch/lib中的所有dll放到libtorch/bin中,然后把libtorch/bin加到环境变量的path中,然后重启VS即可,不重启还是会报错。
(二)检测环境是否配置成功
检测是否配置成功:
#include <iostream>
#include <torch/torch.h>
#include<torch/script.h>
using namespace std;
int main() {
torch::Tensor tensor = torch::rand({ 5,3 });
cout << tensor << endl;
system("pause");
return EXIT_SUCCESS;
}
运行结果:
检测能否使用GPU:
#include <iostream>
#include <torch/torch.h>
#include<torch/script.h>
using namespace std;
int main()
{
if (torch::cuda::is_available())
cout << "支持GPU" << endl;
else
cout << "不支持GPU" << endl;
system("pause");
return 0;
}
运行结果:
(三)用C++调用导出的torchscript模型
调用方法:
- 1、直接调用Python中的C++接口,意思是将图片传入Python中的yolov5进行检测,将结果传回C++,实时性较差。(我没有成功实现)
- 2、可以不用转换为onnx和torchscript等类型,直接调用yolov5模型,但是这个是在C++中配置yolov5,进行训练和检测。(我没有成功实现)
- 3、Python搭建yolov5进行训练,然后将模型转换为onnx,再在C++中调用。(我成功实现了)
(注意:此博客的应用要求较为严格,我也是完全按照他的版本进行搭建配置才能够跑通,中途也出现过一些小问题)
强烈建议遵循博主版本搭建:YOLOv5(版本5.0)、OpenCV(版本4.5.3)。
以下代码为根据此博客内容修改。
#include <torch/script.h>
#include <torch/torch.h>
#include<opencv2/opencv.hpp>
#include <iostream>
using namespace std;
using namespace cv;
// 加载类别名
std::vector<std::string> LoadNames(const std::string& path) {
// 存储类别名的向量
std::vector<std::string> class_names;
std::ifstream infile(path); // 打开指定路径的文件
if (infile.is_open()) { // 如果文件成功打开
std::string line;
while (std::getline(infile, line))
class_names.emplace_back(line); // 将每一行的类别名添加到向量中
infile.close(); // 关闭文件
}
else
std::cerr << "文件打开失败!\n"; // 如果打开文件失败,则输出错误信息
return class_names; // 返回类别名向量
}
// 图像填充函数
std::vector<float> LetterboxImage(const cv::Mat& src, cv::Mat& dst, const cv::Size& out_size) {
// 计算原图和填充后图像的尺寸比例
auto in_h = static_cast<float>(src.rows);
auto in_w = static_cast<float>(src.cols);
float out_h = out_size.height;
float out_w = out_size.width;
float scale = std::min(out_w / in_w, out_h / in_h);
// 计算填充后的图像尺寸
int mid_h = static_cast<int>(in_h * scale);
int mid_w = static_cast<int>(in_w * scale);
// 缩放图像
cv::resize(src, dst, cv::Size(mid_w, mid_h));
// 计算上下左右填充的像素数
int top = (static_cast<int>(out_h) - mid_h) / 2;
int down = (static_cast<int>(out_h) - mid_h + 1) / 2;
int left = (static_cast<int>(out_w) - mid_w) / 2;
int right = (static_cast<int>(out_w) - mid_w + 1) / 2;
// 将图像进行边缘填充
cv::copyMakeBorder(dst, dst, top, down, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
// 返回填充信息
std::vector<float> pad_info{ static_cast<float>(left), static_cast<float>(top), scale };
return pad_info;
}
// 检测结果结构体
struct Detection {
cv::Rect bbox; // 检测框坐标
float score; // 置信度
int class_idx; // 类别索引
};
enum Det {
tl_x = 0,
tl_y = 1,
br_x = 2,
br_y = 3,
score = 4,
class_idx = 5
};
// 将张量转换为检测结果
void Tensor2Detection(const at::TensorAccessor<float, 2>& offset_boxes, const at::TensorAccessor<float, 2>& det,
std::vector<cv::Rect>& offset_box_vec, std::vector<float>& score_vec) {
for (int i = 0; i < offset_boxes.size(0); i++) {
// 将张量中的检测框信息和置信度分别转换为矩形对象和浮点数,并存储至向量中
offset_box_vec.emplace_back(cv::Rect(cv::Point(offset_boxes[i][Det::tl_x], offset_boxes[i][Det::tl_y]),
cv::Point(offset_boxes[i][Det::br_x], offset_boxes[i][Det::br_y])));
score_vec.emplace_back(det[i][Det::score]);
}
}
void ScaleCoordinates(std::vector<Detection>& data, float pad_w, float pad_h, float scale, const cv::Size& img_shape) {
auto clip = [](float n, float lower, float upper) {
return std::max(lower, std::min(n, upper));
};
std::vector<Detection> detections;
for (auto& i : data) {
float x1 = (i.bbox.tl().x - pad_w) / scale; // x padding
float y1 = (i.bbox.tl().y - pad_h) / scale; // y padding
float x2 = (i.bbox.br().x - pad_w) / scale; // x padding
float y2 = (i.bbox.br().y - pad_h) / scale; // y padding
x1 = clip(x1, 0, img_shape.width);
y1 = clip(y1, 0, img_shape.height);
x2 = clip(x2, 0, img_shape.width);
y2 = clip(y2, 0, img_shape.height);
i.bbox = cv::Rect(cv::Point(x1, y1), cv::Point(x2, y2));
}
}
// 矩形格式转换函数
torch::Tensor xywh2xyxy(const torch::Tensor& x) {
auto y = torch::zeros_like(x);
// 转换格式: (center x, center y, width, height) ——> (x1, y1, x2, y2)
y.select(1, Det::tl_x) = x.select(1, 0) - x.select(1, 2).div(2);
y.select(1, Det::tl_y) = x.select(1, 1) - x.select(1, 3).div(2);
y.select(1, Det::br_x) = x.select(1, 0) + x.select(1, 2).div(2);
y.select(1, Det::br_y) = x.select(1, 1) + x.select(1, 3).div(2);
return y;
}
// 后处理函数(返回值为检测结果的二维向量的向量)
std::vector<std::vector<Detection>> PostProcessing(const torch::Tensor& detections, float pad_w, float pad_h, float scale,
const cv::Size& img_shape, float conf_thres, float iou_thres) {
/***
* 结果纬度为
index(0):batch
index(1,2):top-left x/y
index(3,4):bottom-right x/y
index(5):score
index(6):class id
* 13*13*3*(1+4)*80
*/
constexpr int item_attr_size = 5; // 每个检测结果的属性数量,包括top-left x/y、bottom-right x/y、score和class id
int batch_size = detections.size(0); // 批次大小
auto num_classes = detections.size(2) - item_attr_size; // 类别数量,例如COCO数据集有80个类别
auto conf_mask = detections.select(2, 4).ge(conf_thres).unsqueeze(2); // 获取置信度大于阈值的候选框
std::vector<std::vector<Detection>> output; // 存储最终的检测结果
output.reserve(batch_size); // 预分配内存
// 在每一个批次迭代所有图像
for (int batch_i = 0; batch_i < batch_size; batch_i++) {
// 应用约束获取当前图像的滤波检测
auto det = torch::masked_select(detections[batch_i], conf_mask[batch_i]).view({ -1, num_classes + item_attr_size });
// 如果没有检测保留,则跳过并开始处理下一个图像
if (0 == det.size(0))
continue;
// 计算总分 = obj_conf * cls_conf, 类似于 x[:, 5:] *= x[:, 4:5]
det.slice(1, item_attr_size, item_attr_size + num_classes) *= det.select(1, 4).unsqueeze(1);
// 矩形框的格式转换
torch::Tensor box = xywh2xyxy(det.slice(1, 0, 4));
// 在每个结果中获得类的最高分数
std::tuple<torch::Tensor, torch::Tensor> max_classes = torch::max(det.slice(1, item_attr_size, item_attr_size + num_classes), 1);
// 类得分
auto max_conf_score = std::get<0>(max_classes);
// 下标
auto max_conf_index = std::get<1>(max_classes);
max_conf_score = max_conf_score.to(torch::kFloat).unsqueeze(1);
max_conf_index = max_conf_index.to(torch::kFloat).unsqueeze(1);
// 形状: n * 6, top-left x/y (0,1), bottom-right x/y (2,3), score(4), class index(5)
det = torch::cat({ box.slice(1, 0, 4), max_conf_score, max_conf_index }, 1);
// 用于批处理NMS
constexpr int max_wh = 4096;
auto c = det.slice(1, item_attr_size, item_attr_size + 1) * max_wh;
auto offset_box = det.slice(1, 0, 4) + c;
std::vector<cv::Rect> offset_box_vec;
std::vector<float> score_vec;
// 将数据复制回cpu
auto offset_boxes_cpu = offset_box.cpu();
auto det_cpu = det.cpu();
const auto& det_cpu_array = det_cpu.accessor<float, 2>();
// 将张量转换为检测结果向量(即检测框转换为矩形坐标,置信度转换为浮点数)
Tensor2Detection(offset_boxes_cpu.accessor<float, 2>(), det_cpu_array, offset_box_vec, score_vec);
// run NMS
std::vector<int> nms_indices;
cv::dnn::NMSBoxes(offset_box_vec, score_vec, conf_thres, iou_thres, nms_indices);// 使用非极大值抑制去除重叠的框
std::vector<Detection> det_vec;
for (int index : nms_indices) {
Detection t;
const auto& b = det_cpu_array[index];
t.bbox = cv::Rect(cv::Point(b[Det::tl_x], b[Det::tl_y]), cv::Point(b[Det::br_x], b[Det::br_y]));
t.score = det_cpu_array[index][Det::score];
t.class_idx = det_cpu_array[index][Det::class_idx];
det_vec.emplace_back(t);
}
ScaleCoordinates(det_vec, pad_w, pad_h, scale, img_shape);
// 保存当前图像的最终检测
output.emplace_back(det_vec);
} // 批处理迭代结束
return output;
}
void Demo(cv::Mat& img, const std::vector<std::vector<Detection>>& detections, const std::vector<std::string>& class_names, bool label = true) {
if (!detections.empty()) {
for (const auto& detection : detections[0]) {
const auto& box = detection.bbox;
float score = detection.score;
int class_idx = detection.class_idx;
cv::rectangle(img, box, cv::Scalar(0, 0, 255), 2);
if (label) {
std::stringstream ss;
ss << std::fixed << std::setprecision(2) << score;
std::string s = class_names[class_idx] + " " + ss.str();
auto font_face = cv::FONT_HERSHEY_DUPLEX;
auto font_scale = 1.0;
int thickness = 1;
int baseline = 0;
auto s_size = cv::getTextSize(s, font_face, font_scale, thickness, &baseline);
cv::rectangle(img,
cv::Point(box.tl().x, box.tl().y - s_size.height - 5),
cv::Point(box.tl().x + s_size.width, box.tl().y),
cv::Scalar(0, 0, 255), -1);
cv::putText(img, s, cv::Point(box.tl().x, box.tl().y - 5),
font_face, font_scale, cv::Scalar(255, 255, 255), thickness);
}
}
}
}
int main() {
// 读取网络权重
torch::jit::script::Module module = torch::jit::load("best.torchscript.pt");// 读取网络模型的权重文件
torch::DeviceType device_type = torch::kCUDA; // 调用CUDA设备
//torch::DeviceType device_type = torch::kCPU; // 调用CPU设备
module.to(device_type); // 将模型部署到CPU上
module.eval();
// 读取图片(必须是三通道的图像)
cv::Mat img = cv::imread("Sc_10.bmp", -1);// -1:读入完整图片,包括alpha通道
resize(img, img, Size(640, 640)); // 将图像大小调整为640x640像素
// 读取类别
std::vector<std::string> class_names = {"huahen"}; // 加载类别名称
//std::vector<std::string> class_names = LoadNames("huahen.names"); // 从文件中加载类别名称
if (class_names.empty()) { // 如果类别名称为空,则返回-1
system("pause");
return -1;
}
// 设置阈值
float conf_thres = 0.4; // 置信度阈值
float iou_thres = 0.5; // IoU(交并比)阈值
// 推理
torch::NoGradGuard no_grad; // 禁用梯度计算
cv::Mat img_input = img.clone(); // 复制输入图像
std::vector<float> pad_info = LetterboxImage(img_input, img_input, cv::Size(640, 640)); // 将输入图像进行letterbox处理
const float pad_w = pad_info[0]; // letterbox的宽度补齐像素数
const float pad_h = pad_info[1]; // letterbox的高度补齐像素数
const float scale = pad_info[2]; // 缩放比例
cv::cvtColor(img_input, img_input, cv::COLOR_BGR2RGB); // 将图像色彩空间从BGR转换为RGB
img_input.convertTo(img_input, CV_32FC3, 1.0f / 255.0f); // 归一化图像,将像素值缩放到[0,1]范围内
auto tensor_img = torch::from_blob(img_input.data, { 1, img_input.rows, img_input.cols, img_input.channels() }).to(device_type); // 将图像数据转换为张量,并移到CPU设备上
tensor_img = tensor_img.permute({ 0, 3, 1, 2 }).contiguous(); // 调整张量维度顺序,由BHWC变为BCHW
std::vector<torch::jit::IValue> inputs;
inputs.emplace_back(tensor_img); // 将张量添加到输入容器中
torch::jit::IValue output = module.forward(inputs); // 前向传播,获取模型输出结果
auto detections = output.toTuple()->elements()[0].toTensor(); // 提取检测结果
auto result = PostProcessing(detections, pad_w, pad_h, scale, img.size(), conf_thres, iou_thres); // 后处理,得到最终的检测结果
if (true) {
Demo(img, result, class_names); // 可视化检测结果
cv::imshow("Result", img);
cv::waitKey(0); // 等待按键退出
}
return 0;
}