yolov8训练检测部署（二）yolov8_cls python版Tensorrt推理加速

CV-King

于 2024-09-09 16:50:23 发布

阅读量792

点赞数 11

分类专栏： YOLOV8训练推理部署文章标签： YOLO python 计算机视觉 opencv 图像处理分类

本文链接：https://blog.csdn.net/HanWenKing/article/details/142062945

版权

YOLOV8训练推理部署专栏收录该内容

2 篇文章 0 订阅

订阅专栏

文章目录

前言
一、tensort_python 安装
二、yolov8_cls tensorrt推理

前言

TensorRT是由NVIDIA开发的一个深度学习推理（inference）引擎，它能够优化深度学习模型以实现高性能的推理。TensorRT可以接受来自多种深度学习框架（如TensorFlow, PyTorch, Caffe等）的模型，并通过一系列优化步骤来提高模型的推理速度，同时减少模型的内存占用。

一、tensort_python 安装

首先下载tensort：
界面为：
在这里插入图片描述
根据自己所在的环境进行安装，我选择的第一个tar压缩包

下载好之后对压缩包进行解压解压完成之后的界面为：

执行下面代码即可完成tensorrt_python：

conda activate Yolov8
cd /home/build/下载/TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4/TensorRT-8.4.3.1/python
pip install xxx.whl

二、yolov8_cls tensorrt推理

注：下面代码的库一般都放在yolov8目录下的models文件夹中，只要将推理代码放在项目文件中的根目录下即可进行推理

from models import TRTModule  # isort:skip
# 从models模块中导入TRTModule类

import argparse
# 导入argparse模块，用于解析命令行参数

from pathlib import Path
# 导入Path类，用于处理文件路径

import cv2
# 导入cv2模块，用于进行图像处理

import torch
# 导入torch模块，提供张量操作以及与TensorRT模块的交互

from models.utils import blob, path_to_list
# 从models.utils模块中导入blob函数和path_to_list函数，用于图像预处理和路径处理

def main(args: argparse.Namespace) -> None:
    # 定义主函数，接收一个argparse.Namespace对象作为参数，不返回任何值
    device = torch.device(args.device)
    # 根据命令行参数创建一个torch.device对象，用于指定TensorRT模块的运行设备

    Engine = TRTModule(args.engine, device)
    # 实例化TRTModule类，传入engine文件路径和设备信息

    H, W = Engine.inp_info[0].shape[-2:]
    # 从Engine对象的inp_info属性中获取输入图像的高度和宽度

    images = path_to_list(args.imgs)
    # 调用path_to_list函数，将命令行参数中的imgs路径转换为图像路径列表

    print(images)
    # 打印图像路径列表

    save_path = Path(args.out_dir)
    # 创建一个Path对象，表示输出目录的路径

    if not args.show and not save_path.exists():
        save_path.mkdir(parents=True, exist_ok=True)
    # 如果命令行参数中未设置显示结果且输出目录不存在，则创建输出目录

    for image in images:
        # 遍历图像路径列表
        save_image = save_path / image.name
        # 构建输出图像的完整路径

        bgr = cv2.imread(str(image))
        # 使用cv2.imread读取图像，得到BGR格式的图像

        draw = bgr.copy()
        # 复制BGR图像，用于后续绘制结果

        bgr = cv2.resize(bgr, (W, H))
        # 将图像缩放到与模型输入尺寸一致

        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        # 将BGR图像转换为RGB图像

        tensor = blob(rgb, return_seg=False)
        # 调用blob函数处理RGB图像，得到模型输入的张量

        tensor = torch.asarray(tensor, device=device)
        # 将张量移动到指定设备上

        # inference
        data = Engine(tensor)
        # 对图像进行推理，得到模型输出

        print(data)
        # 打印模型输出

        score, cls_id = data[0].max(0)
        # 获取最高分数和对应的类别ID

        score = float(score)
        # 将分数转换为浮点数

        cls_id = int(cls_id)
        # 将类别ID转换为整数

        print(cls_id)
        # 打印类别ID

        text = f'{cls_id}:{score:.3f}'
        # 构建要显示的文本

        (_w, _h), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 1)
        # 获取文本的尺寸和基线

        _y1 = min(10, draw.shape[0])
        # 计算文本框的y坐标

        cv2.rectangle(draw, (10, _y1), (10 + _w, _y1 + _h + _bl), (0, 0, 255), -1)
        # 在图像上绘制文本框

        cv2.putText(draw, text, (10, _y1 + _h), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)
        # 在图像上绘制文本

        if args.show:
            cv2.imshow('result', draw)
            cv2.waitKey(0)
        else:
            cv2.imwrite(str(save_image), draw)
        # 根据命令行参数决定是显示图像还是保存图像

def parse_args() -> argparse.Namespace:
    # 定义解析命令行参数的函数，返回一个argparse.Namespace对象
    parser = argparse.ArgumentParser()
    # 创建一个ArgumentParser对象

    parser.add_argument('--engine', type=str, help='Engine file',default="/home/build/下载/ultralytics-main－２/smoke.engine")
    # 添加命令行参数engine，用于指定engine文件路径

    parser.add_argument('--imgs', type=str, help='Images file',default="/home/build/桌面/date/test")
    # 添加命令行参数imgs