devbox
我家小花儿
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

c++读取yolov5模型进行目标检测（读取摄像头实时监测）_yolov5 c++视频检测

作者：我家小花儿 | 2024-03-24 06:03:50

赞

踩

yolov5 c++视频检测

文章介绍

本文是篇基于yolov5模型的一个工程，主要是利用c++将yolov5模型进行调用并测试，从而实现目标检测任务 任务过程中主要重点有两个，第一 版本问题，第二配置问题

一，所需软件及版本

      训练部分 pytorch==1.13.0  opencv==3.4.1   其他的直接pip即可

      c++部署 

       vs2019或者vs2022    

        libtorch-1.13.0

        opencv==3.4.1    链接：https://pan.baidu.com/s/1XPWUNfS7PTFiDkHTG8yvcQ 
提取码：d9g4

        有的可能需要cmake反正我没用    链接：https://pan.baidu.com/s/1-eLo7ecgQg94Mjtw-pQcXw 
提取码：rg0x

二，安装vs

官网地址：
Visual Studio 较旧的下载 - 2019、2017、2015 和以前的版本

 上诉链接可能为2017推荐安装 Visual Studio Installer  2019或者2022

环境配置以及任务准备可以借鉴我上一篇文章

libtorch-yolov5部署pytorch版本_该醒醒了~的博客-CSDN博客

好的屁话不多说，正文开始

首先在vs中创建新文件

在源文件中新建一个cpp文件，在头文件新建一个.h 头文件

下载yolov5 libtorch

文件链接：https://pan.baidu.com/s/1oIP1btJd10gQddxAHijg7w 
提取码：lntf

  • 粘贴 src/YoloV5.cpp 中的代码到上面的 YoloV5.cpp 文件中
  • 粘贴 nclude/YoloV5.h 中的代码到上面的 YoloV5.h 文件中
  • 更改 YoloV5.cpp 中头文件引入方式为 "YoloV5.h
  • "

改为

 在源文件里新建一个main.cpp 文件 此文件是用来调用yolov5的

将代码复制到main.cpp中

这是读取摄像头实时监测的

  1. #include "YoloV5.h"
  2.  
  3. int main()
  4. {
  5. 	// 第二个参数为是否启用 cuda 详细用法可以参考 YoloV5.h 文件
  6. 	YoloV5 yolo("C:/Users/hwx/Documents/Github/YoloV5-LibTorch/test/yolov5s.cuda.pt", true);
  7. 	// 读取分类标签（我们用的官方的所以这里是 coco 中的分类）
  8. 	// 其实这些代码无所谓哪 只是后面预测出来的框没有标签罢了
  9. 	std::ifstream f("C:/Users/hwx/Documents/Github/YoloV5-LibTorch/test/coco.txt");
  10. 	std::string name = "";
  11. 	int i = 0;
  12. 	std::map<int, std::string> labels;
  13. 	while (std::getline(f, name))
  14. 	{
  15. 		labels.insert(std::pair<int, std::string>(i, name));
  16. 		i++;
  17. 	}
  18. 	// 用 OpenCV 打开摄像头读取文件（你随便咋样获取图片都OK哪）
  19. 	cv::VideoCapture cap = cv::VideoCapture(0);
  20. 	// 设置宽高 无所谓多宽多高后面都会通过一个算法转换为固定宽高的
  21. 	// 固定宽高值应该是你通过YoloV5训练得到的模型所需要的
  22. 	// 传入方式是构造 YoloV5 对象时传入 width 默认值为 640，height 默认值为 640
  23. 	cap.set(cv::CAP_PROP_FRAME_WIDTH, 1000);
  24. 	cap.set(cv::CAP_PROP_FRAME_HEIGHT, 800);
  25. 	cv::Mat frame;
  26. 	while (cap.isOpened())
  27. 	{
  28. 		// 读取一帧
  29. 		cap.read(frame);
  30. 		if (frame.empty())
  31. 		{
  32. 			std::cout << "Read frame failed!" << std::endl;
  33. 			break;
  34. 		}
  35. 		// 预测
  36. 		// 简单吧，两行代码预测结果就出来了，封装的还可以吧 嘚瑟
  37. 		std::vector<torch::Tensor> r = yolo.prediction(frame);
  38. 		// 画框根据你自己的项目调用相应的方法，也可以不画框自己处理
  39. 		frame = yolo.drawRectangle(frame, r[0], labels);
  40. 		// show 图片
  41. 		cv::imshow("", frame);
  42. 		if (cv::waitKey(1) == 27) break;
  43. 	}
  44. 	return 0;
  45. }

这个是读取文件夹内所有的图片

  1. #if 0
  2. #include "YOLOv5.h"
  3. #include"Ex.h"
  4. #include<opencv2\opencv.hpp>
  5. #include<io.h>
  6. #include<iostream>
  7.  
  8. int main()
  9. {
  10. 	YoloV5 yolo("D:\\Besktop\\best.torchscript.pt", true);
  11. 	// 读取分类标签
  12. 	std::ifstream f("D:\\Besktop\\voc.txt");
  13. 	std::string name = "";
  14. 	int i = 0;
  15.  
  16. 	std::map<int, std::string> labels;
  17. 	while (std::getline(f, name))
  18. 	{
  19. 		labels.insert(std::pair<int, std::string>(i, name));
  20. 		std::cout << labels << std::endl;
  21. 		i++;
  22. 	}
  23. 	//cv::Mat frame = cv::imread("D:\\Besktop\\000\\划伤_2023032218553818.bmp");
  24. 	string path = "D:\\Besktop\\000\\";
  25. 	String dest = "D:\\Besktop\\1\\";
  26. 	String savedfilename;
  27. 	int len = path.length();
  28. 	vector<cv::String> filenames;
  29. 	cv::glob(path, filenames);
  30. 	for (int i = 0; i < filenames.size(); i++) 
  31. 	{
  32. 		Mat frame;
  33. 		frame = imread(filenames[i], i);
  34. 		//frame = 255 - frame;   //对每一张图片取反
  35. 		savedfilename = dest + filenames[i].substr(len);
  36. 		cout << savedfilename << endl;
  37. 		// 预测
  38. 		std::vector<torch::Tensor> r = yolo.prediction(frame);
  39. 		std::cout << r << std::endl;
  40. 		// 画框
  41. 		frame = yolo.drawRectangle(frame, r[0], labels);
  42. 		//bool is = yolo.existencePrediction(r);
  43. 		//std::cout << is << std::endl;
  44. 		// show 图片
  45. 		//cv::imshow("", frame);
  46. 		//imwrite(fileName, frame);
  47. 		imwrite(savedfilename, frame);
  48. 		cv::waitKey(0);
  49. 		//if (cv::waitKey(1) == 27);
  50. 		
  51. 	}
  52. 	return 0;
  53. }
  54. #endif // 0

读取一张图片

  1.  
  2. #if 1
  3. #include "YoloV5.h"
  4.  
  5. int main()
  6. {
  7. 	YoloV5 yolo("../dataset/best.torchscript.pt", true);
  8. 	// 读取分类标签
  9. 	std::ifstream f("../dataset/voc.txt");
  10. 	std::string name = "";
  11. 	int i = 0;
  12. 	std::map<int, std::string> labels;
  13. 	while (std::getline(f, name))
  14. 	{
  15. 		labels.insert(std::pair<int, std::string>(i, name));
  16. 		i++;
  17. 	}
  18. 	// 用 OpenCV 打开摄像头读取文件
  19. 	//cv::VideoCapture cap = cv::VideoCapture(0);
  20. 	//cap.set(cv::CAP_PROP_FRAME_WIDTH, 1000);
  21. 	//cap.set(cv::CAP_PROP_FRAME_HEIGHT, 800);
  22. 	//cv::Mat frame;
  23. 	//while (cap.isOpened())
  24. 	//{
  25. 	//	// 读取一帧
  26. 	//	cap.read(frame);
  27. 	//	if (frame.empty())
  28. 	//	{
  29. 	//		std::cout << "Read frame failed!" << std::endl;
  30. 	//		break;
  31. 	//	}
  32. 		cv::Mat frame = cv::imread("D:\\Besktop\\000\\断栅_2.bmp");
  33. 		// 预测
  34. 		std::vector<torch::Tensor> r = yolo.prediction(frame);
  35. 		std::cout << r << std::endl;
  36. 		// 画框处理
  37. 		frame = yolo.drawRectangle(frame, r[0], labels);
  38. 		// show 图片
  39. 		cv::imshow("", frame);
  40. 		cv::waitKey(0);
  41. 		//if (cv::waitKey(1) == 27) break;
  42. 		return 0;
  43. 	}
  44. 	
  45.  
  46. #endif // 1

在 VC++目录/包含目录 中添加头文件

 

在 VC++目录/库目录 中添加 .lib 库 有就添加没有就不添加

在 输入/附加依赖项 中添加 lib 库名称 有就添加，没有就不添加

路径在你libtorch和opencv文件中的lib文件夹内

  1. torch.lib
  2. torch_cuda.lib
  3. torch_cuda_cu.lib
  4. torch_cuda_cpp.lib
  5. torch_cpu.lib
  6. c10_cuda.lib
  7. caffe2_nvrtc.lib
  8. c10.lib
  9. kineto.lib
  10. dnnl.lib
  11. fbgemm.lib
  12. asmjit.lib
  13. XNNPACK.lib
  14. cpuinfo.lib
  15. clog.lib
  16. libprotoc.lib
  17. pthreadpool.lib
  18. libprotobuf.lib
  19. libprotobuf-lite.lib
  20. opencv_world341.lib

注意   将此libtorch和opencv文件夹下的.dll 文件复制到 你项目文件下的环境内

libtorch   .dll文件打开就可以看到

opencv 的藏得比较深  .......opencv3.4.1\opencv\build\x64\vc15\lib

我的项目文件添加路劲为.......active2\x64\Release

 /INCLUDE:"?ignore_this_library_placeholder@@YAHXZ" 

然后修改main函数中的模型路径和下方的标签路径最后运行就好了

最后我们再来核对一下

程序中将会有这些文件其中 只需要看main.cpp tesst.cpp YOLOV5.h Yolov5.cpp

这些文件

main.cpp在上面有这里就不粘贴了

tesst.cpp

  1. #if 0
  2. #include "YOLOv5.h"
  3. #include"Ex.h"
  4. #include<opencv2\opencv.hpp>
  5. #include<io.h>
  6. #include<iostream>
  7.  
  8. int main()
  9. {
  10. 	YoloV5 yolo("D:\\Besktop\\best.torchscript.pt", true);
  11. 	// 读取分类标签
  12. 	std::ifstream f("D:\\Besktop\\voc.txt");
  13. 	std::string name = "";
  14. 	int i = 0;
  15.  
  16. 	std::map<int, std::string> labels;
  17. 	while (std::getline(f, name))
  18. 	{
  19. 		labels.insert(std::pair<int, std::string>(i, name));
  20. 		std::cout << labels << std::endl;
  21. 		i++;
  22. 	}
  23. 	//cv::Mat frame = cv::imread("D:\\Besktop\\000\\划伤_2023032218553818.bmp");
  24. 	string path = "D:\\Besktop\\000\\";
  25. 	String dest = "D:\\Besktop\\1\\";
  26. 	String savedfilename;
  27. 	int len = path.length();
  28. 	vector<cv::String> filenames;
  29. 	cv::glob(path, filenames);
  30. 	for (int i = 0; i < filenames.size(); i++) 
  31. 	{
  32. 		Mat frame;
  33. 		frame = imread(filenames[i], i);
  34. 		//frame = 255 - frame;   //对每一张图片取反
  35. 		savedfilename = dest + filenames[i].substr(len);
  36. 		cout << savedfilename << endl;
  37. 		// 预测
  38. 		std::vector<torch::Tensor> r = yolo.prediction(frame);
  39. 		std::cout << r << std::endl;
  40. 		// 画框
  41. 		frame = yolo.drawRectangle(frame, r[0], labels);
  42. 		//bool is = yolo.existencePrediction(r);
  43. 		//std::cout << is << std::endl;
  44. 		// show 图片
  45. 		//cv::imshow("", frame);
  46. 		//imwrite(fileName, frame);
  47. 		imwrite(savedfilename, frame);
  48. 		cv::waitKey(0);
  49. 		//if (cv::waitKey(1) == 27);
  50. 		
  51. 	}
  52. 	return 0;
  53. }
  54. #endif // 0

yolov5.h

  1. #pragma once
  2. #include <torch/torch.h>
  3. #include <torch/script.h>
  4. #include <opencv2/opencv.hpp>
  5. #include <iostream>
  6. #include <ctime>
  7. /**
  8.  * ImageResizeData 图片处理过后保存图片的数据结构
  9.  */
  10. class ImageResizeData
  11. {
  12. public:
  13. 	// 添加处理过后的图片
  14. 	void setImg(cv::Mat img);
  15. 	// 获取处理过后的图片
  16. 	cv::Mat getImg();
  17. 	// 当原始图片宽高比大于处理过后图片宽高比时此函数返回 true
  18. 	bool isW();
  19. 	// 当原始图片高宽比大于处理过后图片高宽比时此函数返回 true
  20. 	bool isH();
  21. 	// 添加处理之后图片的宽
  22. 	void setWidth(int width);
  23. 	// 获取处理之后图片的宽
  24. 	int getWidth();
  25. 	// 添加处理之后图片的高
  26. 	void setHeight(int height);
  27. 	// 获取处理之后图片的高
  28. 	int getHeight();
  29. 	// 添加原始图片的宽
  30. 	void setW(int w);
  31. 	// 获取原始图片的宽
  32. 	int getW();
  33. 	// 添加原始图片的高
  34. 	void setH(int h);
  35. 	// 获取原始图片的高
  36. 	int getH();
  37. 	// 添加从原始图片到处理过后图片所添加黑边大小
  38. 	void setBorder(int border);
  39. 	// 获取从原始图片到处理过后图片所添加黑边大小
  40. 	int getBorder();
  41. private:
  42. 	// 处理过后图片高
  43. 	int height;
  44. 	// 处理过后图片宽
  45. 	int width;
  46. 	// 原始图片宽
  47. 	int w;
  48. 	// 原始图片高
  49. 	int h;
  50. 	// 从原始图片到处理图片所添加的黑边大小
  51. 	int border;
  52. 	// 处理过后的图片
  53. 	cv::Mat img;
  54. };
  55.  
  56. /**
  57.  * YoloV5 的实现类
  58.  */
  59. class YoloV5
  60. {
  61. public:
  62. 	/**
  63. 	 * 构造函数
  64. 	 * @param ptFile yoloV5 pt文件路径
  65. 	 * @param isCuda 是否使用 cuda 默认不起用
  66. 	 * @param height yoloV5 训练时图片的高
  67. 	 * @param width yoloV5 训练时图片的宽
  68. 	 * @param confThres 非极大值抑制中的 scoreThresh
  69. 	 * @param iouThres 非极大值抑制中的 iouThresh
  70. 	 */
  71. 	YoloV5(std::string ptFile, bool isCuda = false, bool isHalf = false, int height = 640, int width = 640, float confThres = 0.25, float iouThres = 0.45);
  72. 	/**
  73. 	 * 预测函数
  74. 	 * @param data 语言预测的数据格式 (batch, rgb, height, width)
  75. 	 */
  76. 	std::vector<torch::Tensor> prediction(torch::Tensor data);
  77. 	/**
  78. 	 * 预测函数
  79. 	 * @param filePath 需要预测的图片路径
  80. 	 */
  81. 	std::vector<torch::Tensor> prediction(std::string filePath);
  82. 	/**
  83. 	 * 预测函数
  84. 	 * @param img 需要预测的图片
  85. 	 */
  86. 	std::vector<torch::Tensor> prediction(cv::Mat img);
  87. 	/**
  88. 	 * 预测函数
  89. 	 * @param imgs 需要预测的图片集合
  90. 	 */
  91. 	std::vector<torch::Tensor> prediction(std::vector <cv::Mat> imgs);
  92. 	/**
  93. 	 * 改变图片大小的函数
  94. 	 * @param img 原始图片
  95. 	 * @param height 要处理成的图片的高
  96. 	 * @param width 要处理成的图片的宽
  97. 	 * @return 封装好的处理过后图片数据结构
  98. 	 */
  99. 	static ImageResizeData resize(cv::Mat img, int height, int width);
  100. 	/**
  101. 	 * 改变图片大小的函数
  102. 	 * @param img 原始图片
  103. 	 * @return 封装好的处理过后图片数据结构
  104. 	 */
  105. 	ImageResizeData resize(cv::Mat img);
  106. 	/**
  107. 	 * 改变图片大小的函数
  108. 	 * @param imgs 原始图片集合
  109. 	 * @param height 要处理成的图片的高
  110. 	 * @param width 要处理成的图片的宽
  111. 	 * @return 封装好的处理过后图片数据结构
  112. 	 */
  113. 	static std::vector<ImageResizeData> resize(std::vector <cv::Mat> imgs, int height, int width);
  114. 	/**
  115. 	 * 改变图片大小的函数
  116. 	 * @param imgs 原始图片集合
  117. 	 * @return 封装好的处理过后图片数据结构
  118. 	 */
  119. 	std::vector<ImageResizeData> resize(std::vector <cv::Mat> imgs);
  120. 	/**
  121. 	 * 根据输出结果在给定图片中画出框
  122. 	 * @param imgs 原始图片集合
  123. 	 * @param rectangles 通过预测函数处理好的结果
  124. 	 * @param labels 类别标签
  125. 	 * @param thickness 线宽
  126. 	 * @return 画好框的图片
  127. 	 */
  128. 	std::vector<cv::Mat> drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, std::map<int, std::string> labels, int thickness = 2);
  129. 	/**
  130. 	 * 根据输出结果在给定图片中画出框
  131. 	 * @param imgs 原始图片集合
  132. 	 * @param rectangles 通过预测函数处理好的结果
  133. 	 * @param thickness 线宽
  134. 	 * @return 画好框的图片
  135. 	 */
  136. 	std::vector<cv::Mat> drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, int thickness = 2);
  137. 	/**
  138. 	 * 根据输出结果在给定图片中画出框
  139. 	 * @param imgs 原始图片集合
  140. 	 * @param rectangles 通过预测函数处理好的结果
  141. 	 * @param colors 每种类型对应颜色
  142. 	 * @param labels 类别标签
  143. 	 * @return 画好框的图片
  144. 	 */
  145. 	std::vector<cv::Mat> drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, std::map<int, cv::Scalar> colors, std::map<int, std::string> labels, int thickness = 2);
  146. 	/**
  147. 	 * 根据输出结果在给定图片中画出框
  148. 	 * @param img 原始图片
  149. 	 * @param rectangle 通过预测函数处理好的结果
  150. 	 * @param thickness 线宽
  151. 	 * @return 画好框的图片
  152. 	 */
  153. 	cv::Mat	drawRectangle(cv::Mat img, torch::Tensor rectangle, int thickness = 2);
  154. 	/**
  155. 	 * 根据输出结果在给定图片中画出框
  156. 	 * @param img 原始图片
  157. 	 * @param rectangle 通过预测函数处理好的结果
  158. 	 * @param labels 类别标签
  159. 	 * @param thickness 线宽
  160. 	 * @return 画好框的图片
  161. 	 */
  162. 	cv::Mat	drawRectangle(cv::Mat img, torch::Tensor rectangle, std::map<int, std::string> labels, int thickness = 2);
  163. 	/**
  164. 	 * 根据输出结果在给定图片中画出框
  165. 	 * @param img 原始图片
  166. 	 * @param rectangle 通过预测函数处理好的结果
  167. 	 * @param colos 每种类型对应颜色
  168. 	 * @param labels 类别标签
  169. 	 * @param thickness 线宽
  170. 	 * @return 画好框的图片
  171. 	 */
  172. 	cv::Mat	drawRectangle(cv::Mat img, torch::Tensor rectangle, std::map<int, cv::Scalar> colors, std::map<int, std::string> labels, int thickness = 2);
  173. 	/**
  174. 	 * 用于判断给定数据是否存在预测
  175. 	 * @param clazz 通过预测函数处理好的结果
  176. 	 * @return 如果图片中存在给定某一种分类返回 true
  177. 	 */
  178. 	bool existencePrediction(torch::Tensor clazz);
  179. 	/**
  180. 	 * 用于判断给定数据是否存在预测
  181. 	 * @param classs 通过预测函数处理好的结果
  182. 	 * @return 如果图片集合中存在给定某一种分类返回 true
  183. 	 */
  184. 	bool existencePrediction(std::vector<torch::Tensor> classs);
  185.  
  186. private:
  187. 	// 是否启用 cuda
  188. 	bool isCuda;
  189. 	// 是否使用半精度
  190. 	bool isHalf;
  191. 	// 非极大值抑制中的第一步数据清理
  192. 	float confThres;
  193. 	// 非极大值抑制中 iou
  194. 	float iouThres;
  195. 	// 模型所需要的图片的高
  196. 	float height;
  197. 	// 模型所需要的图片的宽
  198. 	float width;
  199. 	// 画框颜色 map
  200. 	std::map<int, cv::Scalar> mainColors;
  201. 	// 模型
  202. 	torch::jit::script::Module model;
  203. 	// 随机获取一种颜色
  204. 	cv::Scalar getRandScalar();
  205. 	// 图片通道转换为 rgb
  206. 	cv::Mat img2RGB(cv::Mat img);
  207. 	// 图片变为 Tensor
  208. 	torch::Tensor img2Tensor(cv::Mat img);
  209. 	// (center_x center_y w h) to (left, top, right, bottom)
  210. 	torch::Tensor xywh2xyxy(torch::Tensor x);
  211. 	// 非极大值抑制算法
  212. 	torch::Tensor nms(torch::Tensor bboxes, torch::Tensor scores, float thresh);
  213. 	// 预测出来的框根据原始图片还原算法
  214. 	std::vector<torch::Tensor> sizeOriginal(std::vector<torch::Tensor> result, std::vector<ImageResizeData> imgRDs);
  215. 	// 非极大值抑制算法整体
  216. 	std::vector<torch::Tensor> non_max_suppression(torch::Tensor preds, float confThres = 0.25, float iouThres = 0.45);
  217. };

 yolov5.cpp

  1. #include "YoloV5.h"
  2.  
  3. YoloV5::YoloV5(std::string ptFile, bool isCuda, bool isHalf, int height, int width, float confThres, float iouThres)
  4. {
  5. 	model = torch::jit::load(ptFile);
  6. 	if (isCuda)
  7. 	{
  8. 		model.to(torch::kCUDA);
  9. 	}
  10. 	if (isHalf)
  11. 	{
  12. 		model.to(torch::kHalf);
  13. 	}
  14. 	this->height = height;
  15. 	this->width = width;
  16. 	this->isCuda = isCuda;
  17. 	this->iouThres = iouThres;
  18. 	this->confThres = confThres;
  19. 	this->isHalf = isHalf;
  20. 	model.eval();
  21. 	unsigned seed = time(0);
  22. 	std::srand(seed);
  23. }
  24.  
  25. std::vector<torch::Tensor> YoloV5::non_max_suppression(torch::Tensor prediction, float confThres, float iouThres)
  26. {
  27. 	torch::Tensor xc = prediction.select(2, 4) > confThres;
  28. 	int maxWh = 4096;
  29. 	int maxNms = 30000;
  30. 	std::vector<torch::Tensor> output;
  31. 	for (int i = 0; i < prediction.size(0); i++)
  32. 	{
  33. 		output.push_back(torch::zeros({ 0, 6 }));
  34. 	}
  35. 	for (int i = 0; i < prediction.size(0); i++)
  36. 	{
  37. 		torch::Tensor x = prediction[i];
  38. 		x = x.index_select(0, torch::nonzero(xc[i]).select(1, 0));
  39. 		if (x.size(0) == 0) continue;
  40.  
  41. 		x.slice(1, 5, x.size(1)).mul_(x.slice(1, 4, 5));
  42. 		torch::Tensor box = xywh2xyxy(x.slice(1, 0, 4));
  43. 		std::tuple<torch::Tensor, torch::Tensor> max_tuple = torch::max(x.slice(1, 5, x.size(1)), 1, true);
  44. 		x = torch::cat({ box, std::get<0>(max_tuple), std::get<1>(max_tuple) }, 1);
  45. 		x = x.index_select(0, torch::nonzero(std::get<0>(max_tuple) > confThres).select(1, 0));
  46. 		int n = x.size(0);
  47. 		if (n == 0)
  48. 		{
  49. 			continue;
  50. 		}
  51. 		else if (n > maxNms)
  52. 		{
  53. 			x = x.index_select(0, x.select(1, 4).argsort(0, true).slice(0, 0, maxNms));
  54. 		}
  55. 		torch::Tensor c = x.slice(1, 5, 6) * maxWh;
  56. 		torch::Tensor boxes = x.slice(1, 0, 4) + c, scores = x.select(1, 4);
  57. 		torch::Tensor ix = nms(boxes, scores, iouThres).to(x.device());
  58. 		output[i] = x.index_select(0, ix).cpu();
  59. 	}
  60. 	return output;
  61. }
  62.  
  63. cv::Scalar YoloV5::getRandScalar()
  64. {
  65. 	return cv::Scalar(std::rand() % 256, std::rand() % 256, std::rand() % 256);
  66. }
  67.  
  68. cv::Mat YoloV5::img2RGB(cv::Mat img)
  69. {
  70. 	int imgC = img.channels();
  71. 	if (imgC == 1)
  72. 	{
  73. 		cv::cvtColor(img, img, cv::COLOR_GRAY2RGB);
  74. 	}
  75. 	else
  76. 	{
  77. 		cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
  78. 	}
  79. 	return img;
  80. }
  81.  
  82. torch::Tensor YoloV5::img2Tensor(cv::Mat img)
  83. {
  84. 	torch::Tensor data = torch::from_blob(img.data, { (int)height, (int)width, 3 }, torch::kByte);
  85. 	data = data.permute({ 2, 0, 1 });
  86. 	data = data.toType(torch::kFloat);
  87. 	data = data.div(255);
  88. 	data = data.unsqueeze(0);
  89. 	return data;
  90. }
  91.  
  92. torch::Tensor YoloV5::xywh2xyxy(torch::Tensor x)
  93. {
  94. 	torch::Tensor y = x.clone();
  95. 	y.select(1, 0) = x.select(1, 0) - x.select(1, 2) / 2;
  96. 	y.select(1, 1) = x.select(1, 1) - x.select(1, 3) / 2;
  97. 	y.select(1, 2) = x.select(1, 0) + x.select(1, 2) / 2;
  98. 	y.select(1, 3) = x.select(1, 1) + x.select(1, 3) / 2;
  99. 	return y;
  100. }
  101.  
  102. torch::Tensor YoloV5::nms(torch::Tensor bboxes, torch::Tensor scores, float thresh)
  103. {
  104. 	auto x1 = bboxes.select(1, 0);
  105. 	auto y1 = bboxes.select(1, 1);
  106. 	auto x2 = bboxes.select(1, 2);
  107. 	auto y2 = bboxes.select(1, 3);
  108. 	auto areas = (x2 - x1) * (y2 - y1);
  109. 	auto tuple_sorted = scores.sort(0, true);
  110. 	auto order = std::get<1>(tuple_sorted);
  111.  
  112. 	std::vector<int> keep;
  113. 	while (order.numel() > 0)
  114. 	{
  115. 		if (order.numel() == 1)
  116. 		{
  117. 			auto i = order.item();
  118. 			keep.push_back(i.toInt());
  119. 			break;
  120. 		}
  121. 		else
  122. 		{
  123. 			auto i = order[0].item();
  124. 			keep.push_back(i.toInt());
  125. 		}
  126.  
  127. 		auto order_mask = order.narrow(0, 1, order.size(-1) - 1);
  128.  
  129. 		auto xx1 = x1.index({ order_mask }).clamp(x1[keep.back()].item().toFloat(), 1e10);
  130. 		auto yy1 = y1.index({ order_mask }).clamp(y1[keep.back()].item().toFloat(), 1e10);
  131. 		auto xx2 = x2.index({ order_mask }).clamp(0, x2[keep.back()].item().toFloat());
  132. 		auto yy2 = y2.index({ order_mask }).clamp(0, y2[keep.back()].item().toFloat());
  133. 		auto inter = (xx2 - xx1).clamp(0, 1e10) * (yy2 - yy1).clamp(0, 1e10);
  134.  
  135. 		auto iou = inter / (areas[keep.back()] + areas.index({ order.narrow(0,1,order.size(-1) - 1) }) - inter);
  136. 		auto idx = (iou <= thresh).nonzero().squeeze();
  137. 		if (idx.numel() == 0)
  138. 		{
  139. 			break;
  140. 		}
  141. 		order = order.index({ idx + 1 });
  142. 	}
  143. 	return torch::tensor(keep);
  144. }
  145.  
  146. std::vector<torch::Tensor> YoloV5::sizeOriginal(std::vector<torch::Tensor> result, std::vector<ImageResizeData> imgRDs)
  147. {
  148. 	std::vector<torch::Tensor> resultOrg;
  149. 	for (int i = 0; i < result.size(); i++)
  150. 	{
  151.  
  152. 		torch::Tensor data = result[i];
  153. 		ImageResizeData imgRD = imgRDs[i];
  154. 		for (int j = 0; j < data.size(0); j++)
  155. 		{
  156. 			torch::Tensor tensor = data.select(0, j);
  157. 			// (left, top, right, bottom)
  158. 			if (imgRD.isW())
  159. 			{
  160. 				tensor[1] -= imgRD.getBorder();
  161. 				tensor[3] -= imgRD.getBorder();
  162. 				tensor[0] *= (float)imgRD.getW() / (float)imgRD.getWidth();
  163. 				tensor[2] *= (float)imgRD.getW() / (float)imgRD.getWidth();
  164. 				tensor[1] *= (float)imgRD.getH() / (float)(imgRD.getHeight() - 2 * imgRD.getBorder());
  165. 				tensor[3] *= (float)imgRD.getH() / (float)(imgRD.getHeight() - 2 * imgRD.getBorder());
  166. 			}
  167. 			else
  168. 			{
  169. 				tensor[0] -= imgRD.getBorder();
  170. 				tensor[2] -= imgRD.getBorder();
  171. 				tensor[1] *= (float)imgRD.getH() / (float)imgRD.getHeight();
  172. 				tensor[3] *= (float)imgRD.getH() / (float)imgRD.getHeight();
  173. 				tensor[0] *= (float)imgRD.getW() / (float)(imgRD.getWidth() - 2 * imgRD.getBorder());
  174. 				tensor[2] *= (float)imgRD.getW() / (float)(imgRD.getWidth() - 2 * imgRD.getBorder());
  175. 			}
  176. 			// 加了黑边之后预测结果可能在黑边上，就会造成结果为负数
  177. 			for (int k = 0; k < 4; k++)
  178. 			{
  179. 				if (tensor[k].item().toFloat() < 0)
  180. 				{
  181. 					tensor[k] = 0;
  182. 				}
  183. 			}
  184. 		}
  185.  
  186. 		resultOrg.push_back(data);
  187. 	}
  188. 	return resultOrg;
  189. }
  190.  
  191. std::vector<torch::Tensor> YoloV5::prediction(torch::Tensor data)
  192. {
  193. 	if (!data.is_cuda() && this->isCuda)
  194. 	{
  195. 		data = data.cuda();
  196. 	}
  197. 	if (data.is_cuda() && !this->isCuda)
  198. 	{
  199. 		data = data.cpu();
  200. 	}
  201. 	if (this->isHalf)
  202. 	{
  203. 		data = data.to(torch::kHalf);
  204. 	}
  205. 	torch::Tensor pred = model.forward({ data }).toTuple()->elements()[0].toTensor();
  206. 	return non_max_suppression(pred, confThres, iouThres);
  207. }
  208.  
  209. std::vector<torch::Tensor> YoloV5::prediction(std::string filePath)
  210. {
  211. 	cv::Mat img = cv::imread(filePath);
  212. 	return prediction(img);
  213. }
  214.  
  215. std::vector<torch::Tensor> YoloV5::prediction(cv::Mat img)
  216. {
  217. 	ImageResizeData imgRD = resize(img);
  218. 	cv::Mat reImg = img2RGB(imgRD.getImg());
  219. 	torch::Tensor data = img2Tensor(reImg);
  220. 	std::vector<torch::Tensor> result = prediction(data);
  221. 	std::vector<ImageResizeData> imgRDs;
  222. 	imgRDs.push_back(imgRD);
  223. 	return sizeOriginal(result, imgRDs);
  224. }
  225.  
  226. std::vector<torch::Tensor> YoloV5::prediction(std::vector<cv::Mat> imgs)
  227. {
  228. 	std::vector<ImageResizeData> imageRDs;
  229. 	std::vector<torch::Tensor> datas;
  230. 	for (int i = 0; i < imgs.size(); i++)
  231. 	{
  232. 		ImageResizeData imgRD = resize(imgs[i]);
  233. 		imageRDs.push_back(imgRD);
  234. 		cv::Mat img = img2RGB(imgRD.getImg());
  235. 		datas.push_back(img2Tensor(img));
  236. 	}
  237. 	torch::Tensor data = torch::cat(datas, 0);
  238. 	std::vector<torch::Tensor> result = prediction(data);
  239. 	return sizeOriginal(result, imageRDs);
  240. }
  241.  
  242. ImageResizeData YoloV5::resize(cv::Mat img, int height, int width)
  243. {
  244. 	ImageResizeData imgResizeData;
  245. 	int w = img.cols, h = img.rows;
  246. 	imgResizeData.setH(h);
  247. 	imgResizeData.setW(w);
  248. 	imgResizeData.setHeight(height);
  249. 	imgResizeData.setWidth(width);
  250. 	bool isW = (float)w / (float)h > (float)width / (float)height;
  251.  
  252. 	cv::resize(img, img, cv::Size(
  253. 		isW ? width : (int)((float)height / (float)h * w),
  254. 		isW ? (int)((float)width / (float)w * h) : height));
  255.  
  256. 	w = img.cols, h = img.rows;
  257. 	if (isW)
  258. 	{
  259. 		imgResizeData.setBorder((height - h) / 2);
  260. 		cv::copyMakeBorder(img, img, (height - h) / 2, height - h - (height - h) / 2, 0, 0, cv::BORDER_CONSTANT);
  261. 	}
  262. 	else
  263. 	{
  264. 		imgResizeData.setBorder((width - w) / 2);
  265. 		cv::copyMakeBorder(img, img, 0, 0, (width - w) / 2, width - w - (width - w) / 2, cv::BORDER_CONSTANT);
  266. 	}
  267. 	imgResizeData.setImg(img);
  268. 	return imgResizeData;
  269. }
  270.  
  271. ImageResizeData YoloV5::resize(cv::Mat img)
  272. {
  273. 	return YoloV5::resize(img, height, width);
  274. }
  275.  
  276. std::vector<ImageResizeData> YoloV5::resize(std::vector<cv::Mat> imgs, int height, int width)
  277. {
  278. 	std::vector<ImageResizeData> imgRDs;
  279. 	for (int i = 0; i < imgs.size(); i++)
  280. 	{
  281. 		imgRDs.push_back(YoloV5::resize(imgs[i], height, width));
  282. 	}
  283. 	return imgRDs;
  284. }
  285.  
  286. std::vector<ImageResizeData> YoloV5::resize(std::vector<cv::Mat> imgs)
  287. {
  288. 	return YoloV5::resize(imgs, height, width);
  289. }
  290.  
  291. std::vector<cv::Mat> YoloV5::drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, std::map<int, std::string> labels, int thickness)
  292. {
  293. 	std::map<int, cv::Scalar> colors;
  294. 	return drawRectangle(imgs, rectangles, colors, labels, thickness);
  295. }
  296.  
  297. std::vector<cv::Mat> YoloV5::drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, int thickness)
  298. {
  299. 	std::map<int, cv::Scalar> colors;
  300. 	std::map<int, std::string> labels;
  301. 	return drawRectangle(imgs, rectangles, colors, labels, thickness);
  302. }
  303.  
  304. std::vector<cv::Mat> YoloV5::drawRectangle(std::vector<cv::Mat> imgs, std::vector<torch::Tensor> rectangles, std::map<int, cv::Scalar> colors, std::map<int, std::string> labels, int thickness)
  305. {
  306. 	std::vector<cv::Mat> results;
  307. 	for (int i = 0; i < imgs.size(); i++)
  308. 	{
  309. 		results.push_back(drawRectangle(imgs[i], rectangles[i], colors, labels, thickness));
  310. 	}
  311. 	return results;
  312. }
  313.  
  314. cv::Mat YoloV5::drawRectangle(cv::Mat img, torch::Tensor rectangle, int thickness)
  315. {
  316. 	std::map<int, cv::Scalar> colors;
  317. 	std::map<int, std::string> labels;
  318. 	return drawRectangle(img, rectangle, colors, labels, thickness);
  319. }
  320.  
  321. cv::Mat YoloV5::drawRectangle(cv::Mat img, torch::Tensor rectangle, std::map<int, std::string> labels, int thickness)
  322. {
  323. 	std::map<int, cv::Scalar> colors;
  324. 	return drawRectangle(img, rectangle, colors, labels, thickness);
  325. }
  326.  
  327. cv::Mat YoloV5::drawRectangle(cv::Mat img, torch::Tensor rectangle, std::map<int, cv::Scalar> colors, std::map<int, std::string> labels, int thickness)
  328. {
  329. 	std::map<int, cv::Scalar>::iterator it;
  330. 	std::map<int, std::string>::iterator labelIt;
  331. 	for (int i = 0; i < rectangle.size(0); i++)
  332. 	{
  333. 		int clazz = rectangle[i][5].item().toInt();
  334. 		it = colors.find(clazz);
  335. 		cv::Scalar color = NULL;
  336. 		if (it == colors.end())
  337. 		{
  338. 			it = mainColors.find(clazz);
  339. 			if (it == mainColors.end())
  340. 			{
  341. 				color = getRandScalar();
  342. 				mainColors.insert(std::pair<int, cv::Scalar>(clazz, color));
  343. 			}
  344. 			else
  345. 			{
  346. 				color = it->second;
  347. 			}
  348. 		}
  349. 		else
  350. 		{
  351. 			color = it->second;
  352. 		}
  353. 		cv::rectangle(img, cv::Point(rectangle[i][0].item().toInt(), rectangle[i][1].item().toInt()), cv::Point(rectangle[i][2].item().toInt(), rectangle[i][3].item().toInt()), color, thickness);
  354. 		labelIt = labels.find(clazz);
  355.  
  356. 		std::ostringstream oss;
  357.  
  358. 		if (labelIt != labels.end())
  359. 		{
  360. 			oss << labelIt->second << " ";
  361. 		}
  362.  
  363. 		oss << rectangle[i][4].item().toFloat();
  364. 		std::string label = oss.str();
  365.  
  366. 		cv::putText(img, label, cv::Point(rectangle[i][0].item().toInt(), rectangle[i][1].item().toInt()), cv::FONT_HERSHEY_PLAIN, 1, color, thickness);
  367. 	}
  368. 	return img;
  369. }
  370.  
  371. bool YoloV5::existencePrediction(torch::Tensor clazz)
  372. {
  373. 	return clazz.size(0) > 0 ? true : false;
  374. }
  375.  
  376. bool YoloV5::existencePrediction(std::vector<torch::Tensor> classs)
  377. {
  378. 	for (int i = 0; i < classs.size(); i++)
  379. 	{
  380. 		if (existencePrediction(classs[i]))
  381. 		{
  382. 			return true;
  383. 		}
  384. 	}
  385. 	return false;
  386. }
  387.  
  388.  
  389. void ImageResizeData::setImg(cv::Mat img)
  390. {
  391. 	this->img = img;
  392. }
  393.  
  394. cv::Mat ImageResizeData::getImg()
  395. {
  396. 	return img;
  397. }
  398.  
  399. bool ImageResizeData::isW()
  400. {
  401. 	return (float)w / (float)h > (float)width / (float)height;
  402. }
  403.  
  404. bool ImageResizeData::isH()
  405. {
  406. 	return (float)h / (float)w > (float)height / (float)width;
  407. }
  408.  
  409. void ImageResizeData::setWidth(int width)
  410. {
  411. 	this->width = width;
  412. }
  413.  
  414. int ImageResizeData::getWidth()
  415. {
  416. 	return width;
  417. }
  418.  
  419. void ImageResizeData::setHeight(int height)
  420. {
  421. 	this->height = height;
  422. }
  423.  
  424. int ImageResizeData::getHeight()
  425. {
  426. 	return height;
  427. }
  428.  
  429. void ImageResizeData::setW(int w)
  430. {
  431. 	this->w = w;
  432. }
  433.  
  434. int ImageResizeData::getW()
  435. {
  436. 	return w;
  437. }
  438.  
  439. void ImageResizeData::setH(int h)
  440. {
  441. 	this->h = h;
  442. }
  443.  
  444. int ImageResizeData::getH()
  445. {
  446. 	return h;
  447. }
  448.  
  449. void ImageResizeData::setBorder(int border)
  450. {
  451. 	this->border = border;
  452. }
  453.  
  454. int ImageResizeData::getBorder()
  455. {
  456. 	return border;
  457. }

 这个效果是读取摄像图进行试试检测博文中有将摄像头替换为图片进行检测的案例

 

看得懂就看，看不懂的评论区问我

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/300339
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号