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网址：PyTorch
https://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;
}
1302
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
