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

Yolov8实例分割Tensorrt部署实战

作者：Monodyee | 2024-02-21 09:26:39

赞

踩

yolov8实例分割

目录

0 引言

1 生成onnx模型

2 onnx转为tensorrt的engine模型

3  Tensorrt推理

3.1 yolov8n-seg分割结果

3.2 yolov8s-seg分割结果

3.3 yolov8m-seg分割结果

3.4 yolov8l-seg分割结果

3.5 yolov8x-seg分割结果

0 引言

        ultralytics在github发布了yolov8模型，可实现快速分类、目标检测与实例分割，采用官方yolov8s-seg.pt效果如下图所示：

​

         本文依旧对其中的实例分割模型进行加速推理实战，开发c++版本的tensorrt推理代码，没有过多的文件依赖，就3个cpp程序文件，不夹带私货，可以算是最简单的推理版本了，直接上链接:Yolov8-instance-seg-tensorrt，本人环境为：cuda10.2、cudnn8.2.4、Tensorrt8.0.1.6、Opencv4.5.4。程序中测试了yolov8[n s m l x]-seg.pt，均能正常使用，代码列表如下

  1. ├── CMakeLists.txt
  2. ├── images
  3. │   ├── bus.jpg
  4. │   └── zidane.jpg
  5. ├── logging.h
  6. ├── main1_onnx2trt.cpp
  7. ├── main2_trt_infer.cpp
  8. ├── models
  9. │   ├── yolov8s-seg.engine
  10. │   └── yolov8s-seg.onnx│   
  11.     ├── yolov8n-seg.engine
  12. │   └── yolov8n-seg.onnx
  13. ├── output.jpg
  14. ├── README.md
  15. └── utils.h

1 生成onnx模型

        yolov8提供了安装方法以及对应的使用代码，在网站下载对应的模型之后，采用如下代码生成所需要的onnx模型

  1. pip install ultralytics
  2. yolo task=segment mode=export model=yolov8[n s m l x]-seg.pt  format=onnx opset=12

        这里注意下，我的onnx版本由于不是最新版，所以opset=12，官方模型默认的opset=17，可以在ultralytics/yolo/configs/default.yaml中找到。

2 onnx转为tensorrt的engine模型

        官方代码提供了直接生成engine的方法，但是我不推荐直接用，原因是生成的engine是跟电脑环境有关的，你换了一个环境之后，比如换到部署服务器上，之前电脑生成的engine就不能用了，除非两个电脑的环境一模一样，不然又要安装所需要的python库，所以我们仅生成onnx模型，然后通过tensorrt的api进行模型转换。

        首先clone我的repo，然后用如下几句话

        1.首先定位到你clone的repo目录下，就是Yolov8-instance-seg-tensorrt目录下
        2.复制 yolov8[n s l m x]-seg.onnx 到 models/目录下

        3.运行下列代码，生成转换与推理的可执行文件-->onnx2trt 、trt_infer

  1. mkdir build
  2. cd build
  3. cmake ..
  4. make  
  5. sudo ./onnx2trt ../models/yolov8s-seg.onnx ../models/yolov8s-seg.engine

        生成onnx2trt的代码如下，可以看到我们主要用的api就是onnxparser

  1. #include <iostream>
  2. #include "logging.h"
  3. #include "NvOnnxParser.h"
  4. #include "NvInfer.h"
  5. #include <fstream>
  6. using namespace nvinfer1;
  7. using namespace nvonnxparser;
  8. static Logger gLogger;
  9. int main(int argc,char** argv) {
  10. 	if (argc < 2) {
  11. 		argv[1] = "../../models/yolov8n-seg.onnx";
  12. 		argv[2] = "../../models/yolov8n-seg.engine";
  13. 	}
  14. 	// 1 onnx解析器
  15. 	IBuilder* builder = createInferBuilder(gLogger);
  16. 	const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
  17. 	INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
  18.  
  19. 	nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
  20.  
  21. 	const char* onnx_filename = argv[1];
  22. 	parser->parseFromFile(onnx_filename, static_cast<int>(Logger::Severity::kWARNING));
  23. 	for (int i = 0; i < parser->getNbErrors(); ++i)
  24. 	{
  25. 		std::cout << parser->getError(i)->desc() << std::endl;
  26. 	}
  27. 	std::cout << "successfully load the onnx model" << std::endl;
  28.  
  29. 	// 2build the engine
  30. 	unsigned int maxBatchSize = 1;
  31. 	builder->setMaxBatchSize(maxBatchSize);
  32. 	IBuilderConfig* config = builder->createBuilderConfig();
  33. 	config->setMaxWorkspaceSize(1 << 20);
  34. 	//config->setMaxWorkspaceSize(128 * (1 << 20));  // 16MB
  35. 	config->setFlag(BuilderFlag::kFP16);
  36. 	ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  37.  
  38. 	// 3serialize Model
  39. 	IHostMemory *gieModelStream = engine->serialize();
  40. 	std::ofstream p(argv[2], std::ios::binary);
  41. 	if (!p)
  42. 	{
  43. 		std::cerr << "could not open plan output file" << std::endl;
  44. 		return -1;
  45. 	}
  46. 	p.write(reinterpret_cast<const char*>(gieModelStream->data()), gieModelStream->size());
  47. 	gieModelStream->destroy();
  48. 	std::cout << "successfully generate the trt engine model" << std::endl;
  49. 	return 0;
  50. }

        通过上述操作，我们就可以得到对应的yolov8[n s m l x]-seg.engine模型。

3  Tensorrt推理

        推理代码和之前的yolov5的实例分割不一样，主要区别在于下图，左边是v5-seg，右边是v8-seg。

         v8中output0输出8400个结果，每个结果的维度是116，而v5是25200个结果，每个结果的维度是117。简述一下v8输出的是啥，116为4+80+32，4为box的cx cy w h，80是每个类的置信度，32是分割需要用到的，和v5的区别在于少了目标的置信度，v5是4+1+80+32，这个1就是是否为目标的置信度。

 推理代码如下：

  1. #include "NvInfer.h"
  2. #include "cuda_runtime_api.h"
  3. #include "NvInferPlugin.h"
  4. #include "logging.h"
  5. #include <opencv2/opencv.hpp>
  6. #include "utils.h"
  7. #include <string>
  8. using namespace nvinfer1;
  9. using namespace cv;
  10.  
  11. // stuff we know about the network and the input/output blobs
  12. static const int INPUT_H = 640;
  13. static const int INPUT_W = 640;
  14. static const int _segWidth = 160;
  15. static const int _segHeight = 160;
  16. static const int _segChannels = 32;
  17. static const int CLASSES = 80;
  18. static const int Num_box = 8400;
  19. static const int OUTPUT_SIZE = Num_box * (CLASSES+4 + _segChannels);//output0
  20. static const int OUTPUT_SIZE1 = _segChannels * _segWidth * _segHeight ;//output1
  21.  
  22.  
  23. static const float CONF_THRESHOLD = 0.1;
  24. static const float NMS_THRESHOLD = 0.5;
  25. static const float MASK_THRESHOLD = 0.5;
  26. const char* INPUT_BLOB_NAME = "images";
  27. const char* OUTPUT_BLOB_NAME = "output0";//detect
  28. const char* OUTPUT_BLOB_NAME1 = "output1";//mask
  29.  
  30.  
  31. struct OutputSeg {
  32. 	int id;             //结果类别id
  33. 	float confidence;   //结果置信度
  34. 	cv::Rect box;       //矩形框
  35. 	cv::Mat boxMask;       //矩形框内mask，节省内存空间和加快速度
  36. };
  37.  
  38. void DrawPred(Mat& img,std:: vector<OutputSeg> result) {
  39. 	//生成随机颜色
  40. 	std::vector<Scalar> color;
  41. 	srand(time(0));
  42. 	for (int i = 0; i < CLASSES; i++) {
  43. 		int b = rand() % 256;
  44. 		int g = rand() % 256;
  45. 		int r = rand() % 256;
  46. 		color.push_back(Scalar(b, g, r));
  47. 	}
  48. 	Mat mask = img.clone();
  49. 	for (int i = 0; i < result.size(); i++) {
  50. 		int left, top;
  51. 		left = result[i].box.x;
  52. 		top = result[i].box.y;
  53. 		int color_num = i;
  54. 		rectangle(img, result[i].box, color[result[i].id], 2, 8);
  55. 		
  56. 		mask(result[i].box).setTo(color[result[i].id], result[i].boxMask);
  57. 		char label[100];
  58. 		sprintf(label, "%d:%.2f", result[i].id, result[i].confidence);
  59.  
  60. 		//std::string label = std::to_string(result[i].id) + ":" + std::to_string(result[i].confidence);
  61. 		int baseLine;
  62. 		Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
  63. 		top = max(top, labelSize.height);
  64. 		putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
  65. 	}
  66. 	
  67. 	addWeighted(img, 0.5, mask, 0.8, 1, img); //将mask加在原图上面
  68.  
  69. 	
  70. }
  71.  
  72.  
  73.  
  74. static Logger gLogger;
  75. void doInference(IExecutionContext& context, float* input, float* output, float* output1, int batchSize)
  76. {
  77.     const ICudaEngine& engine = context.getEngine();
  78.  
  79.     // Pointers to input and output device buffers to pass to engine.
  80.     // Engine requires exactly IEngine::getNbBindings() number of buffers.
  81.     assert(engine.getNbBindings() == 3);
  82.     void* buffers[3];
  83.  
  84.     // In order to bind the buffers, we need to know the names of the input and output tensors.
  85.     // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  86.     const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
  87. 	const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
  88. 	const int outputIndex1 = engine.getBindingIndex(OUTPUT_BLOB_NAME1);
  89.  
  90.     // Create GPU buffers on device
  91.     CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));//
  92. 	CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
  93. 	CHECK(cudaMalloc(&buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float)));
  94. 	// cudaMalloc分配内存 cudaFree释放内存 cudaMemcpy或 cudaMemcpyAsync 在主机和设备之间传输数据
  95. 	// cudaMemcpy cudaMemcpyAsync 显式地阻塞传输 显式地非阻塞传输 
  96.     // Create stream
  97.     cudaStream_t stream;
  98.     CHECK(cudaStreamCreate(&stream));
  99.  
  100.     // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  101.     CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
  102.     context.enqueue(batchSize, buffers, stream, nullptr);
  103. 	CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
  104. 	CHECK(cudaMemcpyAsync(output1, buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float), cudaMemcpyDeviceToHost, stream));
  105. 	cudaStreamSynchronize(stream);
  106.  
  107.     // Release stream and buffers
  108.     cudaStreamDestroy(stream);
  109.     CHECK(cudaFree(buffers[inputIndex]));
  110. 	CHECK(cudaFree(buffers[outputIndex]));
  111. 	CHECK(cudaFree(buffers[outputIndex1]));
  112. }
  113.  
  114.  
  115.  
  116. int main(int argc, char** argv)
  117. {
  118. 	if (argc < 2) {
  119. 		argv[1] = "../models/yolov8n-seg.engine";
  120. 		argv[2] = "../images/bus.jpg";
  121. 	}
  122. 	// create a model using the API directly and serialize it to a stream
  123. 	char* trtModelStream{ nullptr }; //char* trtModelStream==nullptr;  开辟空指针后 要和new配合使用，比如89行 trtModelStream = new char[size]
  124. 	size_t size{ 0 };//与int固定四个字节不同有所不同,size_t的取值range是目标平台下最大可能的数组尺寸,一些平台下size_t的范围小于int的正数范围,又或者大于unsigned int. 使用Int既有可能浪费，又有可能范围不够大。
  125.  
  126. 	std::ifstream file(argv[1], std::ios::binary);
  127. 	if (file.good()) {
  128. 		std::cout << "load engine success" << std::endl;
  129. 		file.seekg(0, file.end);//指向文件的最后地址
  130. 		size = file.tellg();//把文件长度告诉给size
  131. 		//std::cout << "\nfile:" << argv[1] << " size is";
  132. 		//std::cout << size << "";
  133.  
  134. 		file.seekg(0, file.beg);//指回文件的开始地址
  135. 		trtModelStream = new char[size];//开辟一个char 长度是文件的长度
  136. 		assert(trtModelStream);//
  137. 		file.read(trtModelStream, size);//将文件内容传给trtModelStream
  138. 		file.close();//关闭
  139. 	}
  140. 	else {
  141. 		std::cout << "load engine failed" << std::endl;
  142. 		return 1;
  143. 	}
  144.  
  145. 	
  146. 	Mat src = imread(argv[2], 1);
  147. 	if (src.empty()) { std::cout << "image load faild" << std::endl; return 1; }
  148. 	int img_width = src.cols;
  149. 	int img_height = src.rows;
  150. 	std::cout << "宽高：" << img_width << " " << img_height << std::endl;
  151. 	// Subtract mean from image
  152. 	static float data[3 * INPUT_H * INPUT_W];
  153. 	Mat pr_img0, pr_img;
  154. 	std::vector<int> padsize;
  155. 	pr_img = preprocess_img(src, INPUT_H, INPUT_W, padsize);       // Resize
  156. 	int newh = padsize[0], neww = padsize[1], padh = padsize[2], padw = padsize[3];
  157. 	float ratio_h = (float)src.rows / newh;
  158. 	float ratio_w = (float)src.cols / neww;
  159. 	int i = 0;// [1,3,INPUT_H,INPUT_W]
  160. 	//std::cout << "pr_img.step" << pr_img.step << std::endl;
  161. 	for (int row = 0; row < INPUT_H; ++row) {
  162. 		uchar* uc_pixel = pr_img.data + row * pr_img.step;//pr_img.step=widthx3 就是每一行有width个3通道的值
  163. 		for (int col = 0; col < INPUT_W; ++col)
  164. 		{
  165.  
  166. 			data[i] = (float)uc_pixel[2] / 255.0;
  167. 			data[i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
  168. 			data[i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.;
  169. 			uc_pixel += 3;
  170. 			++i;
  171. 		}
  172. 	}
  173.  
  174. 	IRuntime* runtime = createInferRuntime(gLogger);
  175. 	assert(runtime != nullptr);
  176. 	bool didInitPlugins = initLibNvInferPlugins(nullptr, "");
  177. 	ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
  178. 	assert(engine != nullptr);
  179. 	IExecutionContext* context = engine->createExecutionContext();
  180. 	assert(context != nullptr);
  181. 	delete[] trtModelStream;
  182.  
  183. 	// Run inference
  184. 	static float prob[OUTPUT_SIZE];
  185. 	static float prob1[OUTPUT_SIZE1];
  186.  
  187. 	//for (int i = 0; i < 10; i++) {//计算10次的推理速度
  188. 	//       auto start = std::chrono::system_clock::now();
  189. 	//       doInference(*context, data, prob, prob1, 1);
  190. 	//       auto end = std::chrono::system_clock::now();
  191. 	//       std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  192. 	//   }
  193. 	auto start = std::chrono::system_clock::now();
  194. 	doInference(*context, data, prob, prob1, 1);
  195. 	auto end = std::chrono::system_clock::now();
  196. 	std::cout << "推理时间：" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  197.  
  198. 	std::vector<int> classIds;//结果id数组
  199. 	std::vector<float> confidences;//结果每个id对应置信度数组
  200. 	std::vector<cv::Rect> boxes;//每个id矩形框
  201. 	std::vector<cv::Mat> picked_proposals;  //后续计算mask
  202.  
  203. 	
  204.  
  205. 	// 处理box
  206. 	int net_length = CLASSES + 4 + _segChannels;
  207. 	cv::Mat out1 = cv::Mat(net_length, Num_box, CV_32F, prob);
  208.  
  209. 	start = std::chrono::system_clock::now();
  210. 	for (int i = 0; i < Num_box; i++) {
  211. 		//输出是1*net_length*Num_box;所以每个box的属性是每隔Num_box取一个值，共net_length个值
  212. 		cv::Mat scores = out1(Rect(i, 4, 1, CLASSES)).clone();
  213. 		Point classIdPoint;
  214. 		double max_class_socre;
  215. 		minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
  216. 		max_class_socre = (float)max_class_socre;
  217. 		if (max_class_socre >= CONF_THRESHOLD) {
  218. 			cv::Mat temp_proto = out1(Rect(i, 4 + CLASSES, 1, _segChannels)).clone();
  219. 			picked_proposals.push_back(temp_proto.t());
  220. 			float x = (out1.at<float>(0, i) - padw) * ratio_w;  //cx
  221. 			float y = (out1.at<float>(1, i) - padh) * ratio_h;  //cy
  222. 			float w = out1.at<float>(2, i) * ratio_w;  //w
  223. 			float h = out1.at<float>(3, i) * ratio_h;  //h
  224. 			int left = MAX((x - 0.5 * w), 0);
  225. 			int top = MAX((y - 0.5 * h), 0);
  226. 			int width = (int)w;
  227. 			int height = (int)h;
  228. 			if (width <= 0 || height <= 0) { continue; }
  229.  
  230. 			classIds.push_back(classIdPoint.y);
  231. 			confidences.push_back(max_class_socre);
  232. 			boxes.push_back(Rect(left, top, width, height));
  233. 		}
  234.  
  235. 	}
  236. 	//执行非最大抑制以消除具有较低置信度的冗余重叠框（NMS）
  237. 	std::vector<int> nms_result;
  238. 	cv::dnn::NMSBoxes(boxes, confidences, CONF_THRESHOLD, NMS_THRESHOLD, nms_result);
  239. 	std::vector<cv::Mat> temp_mask_proposals;
  240. 	std::vector<OutputSeg> output;
  241.     Rect holeImgRect(0, 0, src.cols, src.rows);
  242. 	for (int i = 0; i < nms_result.size(); ++i) {
  243. 		int idx = nms_result[i];
  244. 		OutputSeg result;
  245. 		result.id = classIds[idx];
  246. 		result.confidence = confidences[idx];
  247. 		result.box = boxes[idx]&holeImgRect;
  248. 		output.push_back(result);
  249. 		temp_mask_proposals.push_back(picked_proposals[idx]);
  250. 	}
  251.  
  252. 	// 处理mask
  253. 	Mat maskProposals;
  254. 	for (int i = 0; i < temp_mask_proposals.size(); ++i)
  255. 		maskProposals.push_back(temp_mask_proposals[i]);
  256.  
  257. 	Mat protos = Mat(_segChannels, _segWidth * _segHeight, CV_32F, prob1);
  258. 	Mat matmulRes = (maskProposals * protos).t();//n*32 32*25600 A*B是以数学运算中矩阵相乘的方式实现的，要求A的列数等于B的行数时
  259. 	Mat masks = matmulRes.reshape(output.size(), { _segWidth,_segHeight });//n*160*160
  260.  
  261. 	std::vector<Mat> maskChannels;
  262. 	cv::split(masks, maskChannels);
  263. 	Rect roi(int((float)padw / INPUT_W * _segWidth), int((float)padh / INPUT_H * _segHeight), int(_segWidth - padw / 2), int(_segHeight - padh / 2));
  264. 	for (int i = 0; i < output.size(); ++i) {
  265. 		Mat dest, mask;
  266. 		cv::exp(-maskChannels[i], dest);//sigmoid
  267. 		dest = 1.0 / (1.0 + dest);//160*160
  268. 		dest = dest(roi);
  269. 		resize(dest, mask, cv::Size(src.cols, src.rows), INTER_NEAREST);
  270. 		//crop----截取box中的mask作为该box对应的mask
  271. 		Rect temp_rect = output[i].box;
  272. 		mask = mask(temp_rect) > MASK_THRESHOLD;
  273. 		output[i].boxMask = mask;
  274. 	}
  275. 	end = std::chrono::system_clock::now();
  276. 	std::cout << "后处理时间：" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  277.  
  278. 	DrawPred(src, output);
  279. 	cv::imshow("output.jpg", src);
  280. 	char c = cv::waitKey(0);
  281. 	
  282. 	// Destroy the engine
  283.     context->destroy();
  284.     engine->destroy();
  285.     runtime->destroy();
  286.  
  287. 	system("pause");
  288.     return 0;
  289. }

 最后各网络的结果如下图所示：

3.1 yolov8n-seg分割结果

3.2 yolov8s-seg分割结果

3.3 yolov8m-seg分割结果

3.4 yolov8l-seg分割结果

3.5 yolov8x-seg分割结果

后处理的时间分为检测框和分割结果的处理。由于都是用opencv自带函数，或者opencv矩阵处理的，消耗的时间较多，所以速度上有很大的提升空间。可以参考wangxinyu大佬的tensorrtx项目中的yolov5-7.0的后处理方式，速度很快

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】
推荐阅读
相关标签

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

  

闽ICP备14008679号