[cmake]基于C++实现yolov5-seg实例分割tensorrt模型部署演示源码

作品简介

 [cmake]基于C++实现yolov5-seg实例分割tensorrt模型部署YOLOv5-Seg-TensorRT

项目概述

本项目实现了基于YOLOv5分割模型的TensorRT加速推理，支持实例分割功能。通过TensorRT的优化，能够高效处理图像和视频输入，实时输出检测结果和分割掩码。

算法原理

1. 模型架构

  • • 基于YOLOv5目标检测模型，扩展了实例分割功能
  • • 使用TensorRT进行模型加速，支持FP16精度
  • • 输入尺寸：640x640
  • • 输出包括：目标检测结果和分割掩码

2.工作流程

  1. 1. 预处理 ：
  • • 调整图像大小并保持宽高比
  • • 填充图像至640x640尺寸
  • • 转换为CHW格式并归一化
  1. 2. 推理 ：
  • • 使用TensorRT引擎执行模型推理
  • • 输出两个结果：检测结果和分割原型
  1. 3. 后处理 ：
  • • 解析检测结果，过滤低置信度目标
  • • 应用非最大抑制（NMS）
  • • 生成分割掩码并调整至原始图像尺寸
  1. 4. 可视化 ：
  • • 绘制边界框和类别标签
  • • 应用分割掩码并与原图混合

安装环境

必需依赖

  • • Visual Studio 2019
  • • CMake 3.30.1
  • • OpenCV 4.9.0
  • • TensorRT 8.6.1.6
  • • CUDA 11.8.0
  • • cuDNN 8.9.7

环境配置

  1. 1. 安装CUDA和cuDNN
  2. 2. 安装TensorRT并设置环境变量
  3. 3. 安装OpenCV并配置路径
  4. 4. 配置CMake以指向正确的依赖路径

运行步骤

1. 构建项目

mkdir build
cd build
cmake ..
cmake --build . --config Release

2. 准备模型

  • • 不可以直接使用提供的 yolov5s-seg.engine 模型文件，需要自己电脑重新转换，因为tensorrt模型依赖于硬件不同硬件模型适配不一样
  • • 或使用ONNX模型文件，程序会自动转换为engine格式


3. 运行推理

命令行参数
yolov5-seg-tensorrt.exe <engine_path> 
示例
# 处理单张图像
yolov5-seg-tensorrt.exe models/yolov5s-seg.engine images/bus.jpg

# 处理视频
yolov5-seg-tensorrt.exe models/yolov5s-seg.engine videos/input.mp4

4. 查看结果

  • • 处理图像时，结果会保存为 原始文件名_output.jpg
  • • 处理视频时，结果会保存为 output.mp4


注意事项

  1. 1. 环境配置 ：确保CUDA、TensorRT和OpenCV的版本与项目要求一致
  2. 2. 模型文件 ：确保模型文件路径正确，支持.onnx和.engine格式
  3. 3. 输入格式 ：支持的图像格式：jpg、jpeg、png
  4. 4. 输入格式 ：支持的视频格式：mp4、avi、m4v、mpeg、mov、mkv、webm
  5. 5. 性能优化 ：首次运行会进行模型预热，之后的推理速度会更快
  6. 6. 内存要求 ：确保GPU有足够的内存来加载模型和处理输入

文件结构与作用

├── src/
│   ├── YOLOv5Seg.cpp        # YOLOv5分割模型实现
│   ├── YOLOv5Seg.h          # YOLOv5分割模型头文件
│   ├── common.h             # 通用工具函数
│   ├── cuda_utils.h         # CUDA相关工具函数
│   ├── logging.h            # 日志处理
│   └── macros.h             # 宏定义
├── main.cpp                 # 主入口文件
├── CMakeLists.txt           # 构建配置文件
├── models/                  # 模型文件目录
│   └── yolov5s-seg.onnx     # ONNX格式模型
├── images/                  # 测试图像目录
│   ├── bus.jpg              # 测试图像1
│   └── zidane.jpg           # 测试图像2
└── build/                   # 构建输出目录

文件作用详解

  1. 1. main.cpp ：
  • • 主入口文件，处理命令行参数
  • • 加载模型和输入文件
  • • 执行预处理、推理和后处理
  • • 保存和显示结果
  1. 2. src/YOLOv5Seg.h ：
  • • 定义YOLOv5Seg类和相关数据结构
  • • 声明模型初始化、预处理、推理和后处理方法
  1. 3. src/YOLOv5Seg.cpp ：
  • • 实现YOLOv5Seg类的所有方法
  • • 处理模型加载、预处理、推理和后处理
  • • 实现分割掩码的生成和处理
  1. 4. src/common.h ：
  • • 通用工具函数，如文件路径处理等
  1. 5. src/cuda_utils.h ：
  • • CUDA相关工具函数，如内存分配和释放
  1. 6. src/logging.h ：
  • • 日志处理功能
  1. 7. src/macros.h ：
  • • 宏定义和常量
  1. 8. CMakeLists.txt ：
  • • 项目构建配置文件
  • • 定义依赖项和编译选项

性能指标

  • • 推理速度 ：取决于GPU性能，在支持的环境下可实现实时处理
  • • 准确率 ：基于YOLOv5模型，保持了原始模型的检测和分割精度
  • • 内存占用 ：模型大小适中，适合在大多数现代GPU上运行

扩展与定制

  1. 1. 模型替换 ：可以使用不同版本的YOLOv5分割模型，只需将ONNX文件替换到models目录
  2. 2. 参数调整 ：可以在YOLOv5Seg.h中调整置信度阈值、NMS阈值等参数
  3. 3. 功能扩展 ：可以添加批量处理、多GPU支持等功能

故障排除

  1. 1. 模型加载失败 ：检查模型文件路径是否正确，确保模型格式正确
  2. 2. 推理错误 ：检查CUDA和TensorRT版本是否兼容
  3. 3. 内存不足 ：减小输入尺寸或使用更小的模型
  4. 4. 性能问题 ：确保启用了FP16精度，使用合适的GPU

 

创作时间：