当前位置：首页 > news >正文

YOLOv5-7.0实例分割+TensorRT部署

news 2025/8/3 21:35:43

一：介绍

将YOLOv5结合分割任务并进行TensorRT部署，是一项既具有挑战性又令人兴奋的任务。分割（Segmentation）任务要求模型不仅能够检测出目标的存在，还要精确地理解目标的边界和轮廓，为每个像素分配相应的类别标签，使得计算机能够对图像进行更深入的理解和解释。而TensorRT作为一种高性能的深度学习推理引擎，能够显著加速模型的推理过程，为实时应用提供了强大的支持。

在本文中，我们将探讨如何将YOLOv5与分割任务相结合，实现同时进行目标检测和像素级别的语义分割。我们将详细介绍模型融合的技术和步骤，并深入讨论如何利用TensorRT对模型进行优化，以实现在嵌入式设备和边缘计算环境中的高效部署。通过阐述实验结果和性能指标，我们将展示这一方法的有效性和潜力，为读者带来关于结合YOLOv5、分割任务和TensorRT部署的全面认识。

二：python

打开pycharm，终端输入pip install labelme
为了方便我们之后标注的工作，需要打开C盘->用户->用户名->.labelmerc文件打开之后，将第一行的auto_save改为true,方便标框需要将create_polygon改为W,方便修改标注框将edit_polygon改为J
下载结束后，pycharm终端输入labelme，打开你数据集的文件夹，进行标注即可，这里不放图片演示了

标注完之后我们需要将json文件转换为txt文件，下面放上所需要的代码

import json
import os
import argparse
from tqdm import tqdmdef convert_label_json(json_dir, save_dir, classes):json_paths = os.listdir(json_dir)classes = classes.split(',')for json_path in tqdm(json_paths):# for json_path in json_paths:path = os.path.join(json_dir, json_path)with open(path, 'r') as load_f:json_dict = json.load(load_f)h, w = json_dict['imageHeight'], json_dict['imageWidth']# save txt pathtxt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))txt_file = open(txt_path, 'w')for shape_dict in json_dict['shapes']:label = shape_dict['label']label_index = classes.index(label)points = shape_dict['points']points_nor_list = []for point in points:points_nor_list.append(point[0] / w)points_nor_list.append(point[1] / h)points_nor_list = list(map(lambda x: str(x), points_nor_list))points_nor_str = ' '.join(points_nor_list)label_str = str(label_index) + ' ' + points_nor_str + '\n'txt_file.writelines(label_str)if __name__ == "__main__":"""python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_ringsts --classes "cat,dogs""""parser = argparse.ArgumentParser(description='json convert to txt params')parser.add_argument('--json-dir', type=str, default=r'json', help='json path dir')parser.add_argument('--save-dir', type=str, default=r'txt',help='txt save dir')parser.add_argument('--classes', type=str,default="1", help='classes')args = parser.parse_args()json_dir = args.json_dirsave_dir = args.save_dirclasses = args.classesconvert_label_json(json_dir, save_dir, classes)

转换为txt文件后，划分一下数据集，进行训练（此步骤有手就行，在此不演示了）

将你训练得到的best.pt通过gen_wts.py转换为wts文件，为了方便操作，将best.pt放入目录下，终端输入：python gen_wts.py -w best.pt

gen_wts.py的代码如下

import sys
import argparse
import os
import struct
import torch
from utils.torch_utils import select_devicedef parse_args():parser = argparse.ArgumentParser(description='Convert .pt file to .wts')parser.add_argument('-w', '--weights', required=True,help='Input weights (.pt) file path (required)')parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)')parser.add_argument('-t', '--type', type=str, default='detect', choices=['detect', 'cls'],help='determines the model is detection/classification')args = parser.parse_args()if not os.path.isfile(args.weights):raise SystemExit('Invalid input file')if not args.output:args.output = os.path.splitext(args.weights)[0] + '.wts'elif os.path.isdir(args.output):args.output = os.path.join(args.output,os.path.splitext(os.path.basename(args.weights))[0] + '.wts')return args.weights, args.output, args.typept_file, wts_file, m_type = parse_args()
print(f'Generating .wts for {m_type} model')# Initialize
device = select_device('cpu')
# Load model
print(f'Loading {pt_file}')
model = torch.load(pt_file, map_location=device)  # load to FP32
model = model['ema' if model.get('ema') else 'model'].float()if m_type == "detect":# update anchor_grid infoanchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]# model.model[-1].anchor_grid = anchor_griddelattr(model.model[-1], 'anchor_grid')  # model.model[-1] is detect layer# The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.model.model[-1].register_buffer("anchor_grid", anchor_grid)model.model[-1].register_buffer("strides", model.model[-1].stride)model.to(device).eval()print(f'Writing into {wts_file}')
with open(wts_file, 'w') as f:f.write('{}\n'.format(len(model.state_dict().keys())))for k, v in model.state_dict().items():vr = v.reshape(-1).cpu().numpy()f.write('{} {} '.format(k, len(vr)))for vv in vr:f.write(' ')f.write(struct.pack('>f', float(vv)).hex())f.write('\n')

三、TensorRT

下载与YOLOv5-7.0对应的tensorrt分割版本wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT network definition API (github.com)
使用cmake解压，嫌麻烦直接自己配置也行
TensorRT的配置我之前文章里面有写，不清楚的可以去看一下Windows YOLOv5-TensorRT部署_tensorrt在windows部署_Mr Dinosaur的博客-CSDN博客
打开config.h，修改一下自己的检测类别和图片大小
打开yolov5_seg.cpp，找到主函数进行文件路径修改
如果你显源码运行麻烦（我就是），当然也可以自行修改去生成它的engine引擎文件，引擎文件生成后即可进行分割测试
分割结果推理速度一般吧，比检测要慢一些