YOLOv7数据标注、环境配置、训练、预测、trt加速

YOLOv7数据标注、环境配置、训练、预测、trt加速

本文的训练教程使用的是2023-07-20从github仓库拉取的代码以及模型,tag是0.1版本,其他版本的训练可能会有所改动
本文主要是教工作室同学使用yolov7,欢迎大佬指出错误

1、下载代码和模型

1.1 下载代码

1.1.1使用git拉取

直接在终端执行,拉取的项目文件会存放在当前的目录里面,拉取之前注意需要先进入自己存放项目的目录

使用这种方法需要提前安装好git

git clone https://github.com/WongKinYiu/yolov7.git

1.1.2下载zip压缩包

打开仓库下载,仓库链接WongKinYiu/yolov7: Implementation of paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors (github.com)

YOLOv7数据标注、环境配置、训练、预测、trt加速_第1张图片

下载之后解压到自己的目录即可

使用zip的方式下载和使用git clone的方式下载得到的文件夹里面的内容是一样的,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第2张图片

1.2 下载预训练模型

在下载下来的文件夹里面创建一个weights的文件夹,用于存放下载下来的模型文件,也可以直接存放在项目根目录下面,看个人喜好

在github仓库里面向下拉找到Performance,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第3张图片

两个位置下载的模型都是一样的,一般使用的是第一个模型,直接下载到刚才创建的文件夹即可

YOLOv7数据标注、环境配置、训练、预测、trt加速_第4张图片

2、环境准备

2.1 、python环境准备

已经有python环境的可以直接跳过

2.1.1、使用anaconda安装python环境

下载地址 Free Download | Anaconda

下载之后直接安装即可,选择仅为个人用户安装,然后后面需要勾选的全部勾选上,如果没勾可能需要单独配置环境变量

单独配置环境变量

使用windows搜索env,进入编辑账户环境变量,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第5张图片

添加如下环境变量

在系统环境变量path里面添加YOLOv7数据标注、环境配置、训练、预测、trt加速_第6张图片

添加到用户path也可以

安装完成之后打开命令提示符

使用conda info查看conda的环境变量是否配置完成,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第7张图片

2.1.2、使用miniconda配置python环境

下载地址 Miniconda — conda documentation

这个的安装和anaconda一样,对照使用

2.2、使用GPU训练

目前只有Nvidia显卡支持使用GPU训练,需要安装cuda和cudnn,使用AMD和intel的显卡的电脑目前只能用CPU训练(可能有大佬能使用这两家的显卡进行训练,我只是小白,大佬勿喷)

2.2.1、 安装cuda和cudnn

  1. cuda

下载地址:

CUDA Toolkit 12.2 Downloads | NVIDIA Developer

下载地址进去默认是最新版,可以拉大最后选择自己需要的版本

YOLOv7数据标注、环境配置、训练、预测、trt加速_第8张图片

版本选择界面是下面这样的

YOLOv7数据标注、环境配置、训练、预测、trt加速_第9张图片

下载下来的cuda安装包是一个.exe文件,直接双击运行,安装过程中可以选择一下安装的位置,其他有需要勾选的地方全部勾选上

默认安装位置是(我使用的是11.3的版本)

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3

配置环境变量

使用windows搜索env,进入编辑账户环境变量,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第10张图片

添加如下环境变量

在系统环境变量path里面添加

image-20230720160247242

注意将这两个路径换成自己的安装路径

然后打开cmd命令框,使用nvcc -Vnvidia-smi分别查看,得到类似下图的即可

YOLOv7数据标注、环境配置、训练、预测、trt加速_第11张图片

  1. cudnn

    下载地址

    cuDNN Download | NVIDIA Developer

YOLOv7数据标注、环境配置、训练、预测、trt加速_第12张图片

cudnn下载下来之后解压到cuda的安装路径即可

YOLOv7数据标注、环境配置、训练、预测、trt加速_第13张图片

解压到这个目录下面即可,如果提示有重复文件选择覆盖

2.2.2使用conda配置yolo环境

创建一个conda环境

直接在终端运行conda create -n yolo python=3.10 在创建环境的同时安装python3.10,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第14张图片

创建完成会有如下输出

YOLOv7数据标注、环境配置、训练、预测、trt加速_第15张图片

使用conda activate yolo 进入创建好的环境

YOLOv7数据标注、环境配置、训练、预测、trt加速_第16张图片

切换到yolov7文件存放的位置

YOLOv7数据标注、环境配置、训练、预测、trt加速_第17张图片

注意里面的文件必须有 requirements.txt

然后配置yolo需要的python环境包,运行pip install -i https://pypi.mirrors.ustc.edu.cn/simple/ -r requirements.txt等待所有包安装完成即可

检测torch能不能调用cuda

在终端进入python环境然后运行

import torch
torch.cuda.is_available()

YOLOv7数据标注、环境配置、训练、预测、trt加速_第18张图片

返回true即可,如果返回的是其他值则需要重新安装torch

需要先卸载之后再安装

卸载指令:pip uninstall torch

进入pytorch官网,链接:PyTorchYOLOv7数据标注、环境配置、训练、预测、trt加速_第19张图片

按照自己的环境配置选择,并复制下方指令,在其后面加入清华源链接安装

安装完成之后再次测试,结果如下即可

YOLOv7数据标注、环境配置、训练、预测、trt加速_第20张图片

到此,环境配置完成

2.3、使用cpu训练

使用cpu训练不需要配置cuda和cudnn,只需配置python环境即可

2.3.1、使用conda配置yolo环境(cpu)

创建一个conda环境

直接在终端运行conda create -n yolo python=3.10 在创建环境的同时安装python3.10,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第21张图片

创建完成会有如下输出

YOLOv7数据标注、环境配置、训练、预测、trt加速_第22张图片

使用conda activate yolo 进入创建好的环境

YOLOv7数据标注、环境配置、训练、预测、trt加速_第23张图片

切换到yolov7文件存放的位置

YOLOv7数据标注、环境配置、训练、预测、trt加速_第24张图片

注意里面的文件必须有 requirements.txt

然后配置yolo需要的python环境包,运行pip install -i https://pypi.mirrors.ustc.edu.cn/simple/ -r requirements.txt等待所有包安装完成即可

3、训练

3.1、数据集标注

使用labelimg,这个组件可以使用pip安装,安装指令

pip install -i https://pypi.mirrors.ustc.edu.cn/simple/ labelimg

安装之后直接使用labelimg运行即可

YOLOv7数据标注、环境配置、训练、预测、trt加速_第25张图片

标记完成之后会生成对应文件的.txt文件,这个保存的是标签数据

还有一个classes.txt,这里保存的是每个标签的顺序,这个顺序在多人标注的时候不能乱,在训练的时候需要用到这个顺序

3.2、调整数据集及配置训练参数

标注完成之后将图片和标注数据分别存放在两个文件夹里面,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第26张图片

将这个数据集存放于yolov7文件夹里面的data里面,类似于

YOLOv7数据标注、环境配置、训练、预测、trt加速_第27张图片

3.2.1、划分数据集

将数据集划分为traintestval三部分,分别用于训练、测试和验证三部分,在数据集较小的时候可以不用测试部分

使用下面代码进行划分,这个代码只划分了训练和验证,没有测试,注意修改35和36行代码,还有21、26行(生成txt文件的路径)

# 划分数据集
import os
import random


def split_files(folder_path, ratio):
    # 获取文件夹中的文件列表
    file_list = os.listdir(folder_path)

    # 计算划分的文件数量
    total_files = len(file_list)
    num_files_part1 = int(total_files * ratio)
    num_files_part2 = total_files - num_files_part1

    # 随机划分文件
    random.shuffle(file_list)
    part1_files = file_list[:num_files_part1]
    part2_files = file_list[num_files_part1:]

    # 将划分结果写入 txt 文件
    with open("./data/coffee/train_list.txt", "w") as f:
        for filename in part1_files:
            abs_path = os.path.join(folder_path, filename)
            f.write(abs_path + "\n")

    with open("./data/coffee/val_list.txt", "w") as f:
        for filename in part2_files:
            abs_path = os.path.join(folder_path, filename)
            f.write(abs_path + "\n")

    print("Splitting files completed.")


if __name__ == "__main__":
    folder_path = r"D:\study_documents\test_file\yolov7\data\coffee\images"  # 替换为实际的文件夹路径
    ratio = 0.9  # 指定训练部分的文件数量比例
    
    split_files(folder_path, ratio)

这个代码划分之后会生成两个txt文件,里面的内容如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第28张图片

3.2.2、修改参数并训练

将data里面的coco.yaml文件修改为如下格式

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: D:\study_documents\test_file\yolov7\data\coffee\train_list.txt
val: D:\study_documents\test_file\yolov7\data\coffee\val_list.txt
# number of classes
nc: 1

# class names
names: ['impurity']

其中nc是标注数据集时候的种类

names是标注时候的类别名字,这个需要和上面提到的classes.txt顺序一样

第二行和第三行按照自己的路径修改即可

修改train.py文件

找到配置参数,我拉取的版本在527行开始,我将需要注意的参数标注出来,可以参考一下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第29张图片

这些参数修改之后直接运行即可,显存小的或者使用cpu训练的注意将**–batch-size**改小一点

3.2.3、对参数的部分理解(个人理解,仅供参考)

  1. --weights(类型:字符串,默认值:‘yolov7.pt’):

    • 作用:指定训练或测试时的初始模型权重路径。
    • 修改:将默认值更改为您所需的初始模型权重路径。
  2. --cfg(类型:字符串,默认值:‘cfg/training/yolov7.yaml’):

    • 作用:指定模型配置文件(model.yaml)的路径,其中定义了YOLOv7模型的架构。
    • 修改:如果您有自定义的模型配置文件,请将默认值更改为其路径。
  3. --data(类型:字符串,默认值:‘data/coco.yaml’):

    • 作用:指定数据配置文件(data.yaml)的路径,其中包含关于数据集的信息,如类别数目、训练和测试图像目录等。
    • 修改:如果您有自定义的数据配置文件,请将默认值更改为其路径。
  4. --hyp(类型:字符串,默认值:‘data/hyp.scratch.p5.yaml’):

    • 作用:指定超参数配置文件(hyp.scratch.p5.yaml)的路径,其中包含了训练的各种超参数。
    • 修改:如果您有自定义的超参数配置文件,请将默认值更改为其路径。
  5. --epochs(类型:整数,默认值:300):

    • 作用:指定训练的总轮数(epochs)。
    • 修改:将默认值更改为所需的训练轮数。
  6. --batch-size(类型:整数,默认值:8):

    • 作用:指定每个GPU上的批次大小。
    • 修改:根据可用的GPU内存,将默认值更改为所需的批次大小。
  7. --img-size(nargs=‘+’, 类型:整数列表,默认值:[640, 640]):

    • 作用:指定训练和测试时的图像大小。
    • 修改:将默认值更改为两个整数的列表,表示所需的图像大小(宽度,高度)。
  8. --rect(action=‘store_true’):

    • 作用:如果提供该参数,则启用矩形训练。矩形训练不会将图像调整为正方形,而是保持原始的长宽比。
  9. --resume(nargs=‘?’, const=True, 默认值:False):

    • 作用:如果提供该参数,则允许从最近的检查点继续训练。
    • 修改:要恢复训练,请将检查点文件的路径作为参数传递。例如:--resume checkpoints/last.pt
  10. --nosave(action=‘store_true’):

    • 作用:如果提供该参数,则只保存最终的检查点,而不保存训练过程中的中间检查点。
  11. --notest(action=‘store_true’):

    • 作用:如果提供该参数,则只在训练结束后进行一次测试,而不会在每个训练轮次后进行测试。
  12. --noautoanchor(action=‘store_true’):

    • 作用:如果提供该参数,则禁用自动锚框检测,不会根据数据集自动调整锚框的尺寸。
  13. --evolve(action=‘store_true’):

    • 作用:如果提供该参数,则启用超参数进化,它会在训练过程中自动搜索更好的超参数。
  14. --bucket(类型:字符串,默认值:‘’):

    • 作用:如果使用Google Cloud Storage,指定用于存储日志和检查点的gsutil存储桶。
  15. --cache-images(action=‘store_true’):

    • 作用:如果提供该参数,则在内存中缓存图像以加快训练速度。
  16. --image-weights(action=‘store_true’):

    • 作用:如果提供该参数,则使用加权图像选择进行训练。
  17. --device(默认值=‘0’):

    • 作用:指定用于训练的CUDA设备(例如,'0’表示使用GPU 0,'0,1’表示使用多个GPU)。
    • 修改:将默认值更改为所需的CUDA设备或如果要使用CPU进行训练,则设置为’cpu’。
  18. --multi-scale(action=‘store_true’):

    • 作用:如果提供该参数,则在训练过程中随机变化图像大小,范围为+/- 50%。
  19. --single-cls(action=‘store_true’):

    • 作用:如果提供该参数,则将多类数据视为单类数据进行训练。
    • 修改:如果有多类数据集,请保留此参数或从代码中删除它。只有在将多类数据视为单类数据时使用此选项。
  20. --adam(action=‘store_true’):

    • 作用:如果提供该参数,则使用Adam优化器进行训练。
    • 修改:如果想要使用其他优化器,请从代码中删除此参数,并根据需要设置优化器。
  21. --sync-bn(action=‘store_true’):

    • 作用:如果提供该参数,则使用SyncBatchNorm,这仅在DDP(Distributed Data Parallel)模式下可用。
  22. --local_rank(类型:整数,默认值:-1):

    • 作用:DDP参数,除非使用分布式数据并行训练,否则不要修改。
  23. --workers(类型:整数,默认值:1):

    • 作用:指定用于加载数据的最大数据加载器(dataloader)工作进程数。
  24. --project(默认值=‘runs/train’):

    • 作用:指定使用Weights & Biases(W&B)库保存训练日志和检查点的路径。

3.3、预测与测试

训练得到的模型保存在runs/train/exp/weights里面,使用best.pt模型,其余模型可以删除,

修改detect.py,修改以下参数

YOLOv7数据标注、环境配置、训练、预测、trt加速_第30张图片

只需修改对应的参数即可

3.3.1、对测试参数的部分理解

  1. --weights:该参数用于指定模型权重(model.pt)文件的路径。可以接受多个路径,用于使用集成模型。修改成自己的模型路径。

  2. --source:该参数用于指定推理的输入数据来源。可以是文件/文件夹路径,也可以是0,表示使用摄像头作为输入源。

  3. --img-size:该参数设置推理时输入图像的尺寸,以像素为单位。默认值为640。

  4. --conf-thres:该参数设置对象置信度阈值,低于该阈值的任何检测到的对象将被丢弃。默认值为0.25。

  5. --iou-thres:该参数设置非最大抑制(NMS)的交并比(IOU)阈值。用于抑制重复的检测结果。默认值为0.45。

  6. --device:该参数用于指定推理所使用的设备,可以选择’cuda’加上GPU设备编号(例如,‘0’或’0,1,2,3’)或’cpu’。

  7. --view-img:这是一个标志参数。如果出现,则会显示推理结果,允许可视化检测到的对象,添加一个default=Ture/Faulse来指定这个参数是否生效。

  8. --save-txt:这是一个标志参数。如果出现,则会将推理结果保存为.txt文件,添加一个default=Ture/Faulse来指定这个参数是否生效。

  9. --save-conf:这是一个标志参数。如果出现,则会将检测到的对象的置信度与.txt结果一起保存,添加一个default=Ture/Faulse来指定这个参数是否生效。

  10. --nosave:这是一个标志参数。如果出现,则会禁止保存带有绘制边界框的图像或视频,用于检测到的对象,添加一个default=Ture/Faulse来指定这个参数是否生效。

  11. --classes:该参数允许您基于特定类别过滤检测到的对象。您可以指定类别ID,只在检测结果中包括这些类别。

  12. --agnostic-nms:这是一个标志参数。如果出现,则会启用无类别非最大抑制,意味着在NMS过程中不考虑类别信息。

  13. --augment:这是一个标志参数。如果出现,则会启用增强推理,即在推理过程中应用数据增强,以提高检测准确性。

  14. --update:这是一个标志参数。如果出现,则会更新--weights参数指定的所有模型。

  15. --project:该参数用于指定保存结果的项目目录。默认值为’runs/detect’。

  16. --name:该参数用于指定实验的名称。默认值为’exp’。

  17. --exist-ok:这是一个标志参数。如果出现,则允许在不递增的情况下使用具有相同名称的现有项目。

  18. --no-trace:这是一个标志参数。如果出现,则会阻止在推理过程中对模型进行追踪。追踪是一种用于优化模型以实现更快推理的技术。

4、TensorRT部署

4.1、TensorRT安装

下载路径

NVIDIA TensorRT 8.x Download | NVIDIA Developer

注意版本

YOLOv7数据标注、环境配置、训练、预测、trt加速_第31张图片

4.1.1、linux配置

下载之后先解压

tar -zxvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz

配置环境变量

# 修改单个用户的环境变量文件
vim ~/.bashrc
# 添加下面两条变量,注意修改成自己的路径
export LD_LIBRARY_PATH=/home/bigdata/mmd/file/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=/home/bigdata/mmd/file/TensorRT-8.6.1.6/lib::$LIBRARY_PATH

# 重新加载环境变量
source ~/.bashrc

然后安装对应的python包

image-20230801144840418

需要安装的python包分别在pythonuufonnx_graphsurgeon里面,python里面需要安装的有三个,选择对应的python版本安装就行

安装指令

pip install *.whl

4.1.2、windows配置

与linux一样的下载之后先解压

将解压目录添加到环境变量的path里面,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第32张图片

只需在系统变量和用户变量选一个添加即可。

接下来复制文件

  • 将 TensorRT-8.6.1.6\include中头文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\include

  • 将TensorRT-8.6.1.6\lib 中所有lib文件复制 到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\lib\x64

  • 将TensorRT-8.6.1.6\lib 中所有dll文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\bin

注意将上面的路径改成自己的安装路径

然后安装相应的python包

这一步与Linux版本的一样,需要安装的python包分别在pythonuufonnx_graphsurgeon里面,python里面需要安装的有三个,选择对应的python版本安装就行

安装指令

pip install *.whl

至此TensotRT已经配置完成。

4.2、模型转换

4.2.1、原生方法转engine

  1. .pt转**.onnx**

修改export.py,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第33张图片

修改自己需要转的模型路径之后直接运行即可,会在对应的位置生成**.mlmodel**、.onnx.torchscript.pttorcscript.ptl结尾的四个类型的文件,我们只需要注意**.onnx**结尾的这个文件即可

将**.onnx结尾的文件复制到下载解压的TensorRT-8.6.1.6\bin**里面,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第34张图片

在这个目录进入终端然后进入运行yolo的环境,如下

image-20230803140109424

使用TensorRT自带的trtexec.exe将**.onnx转成.engine**,转换指令如下

trtexec --onnx=./best.onnx --saveEngine=./best_fp16.engine --fp16 --workspace=200

#Linux同样可以使用这种方法转换,下面是命令

./trtexec --onnx=./best.onnx --saveEngine=./best_fp16.engine --fp16 --workspace=200

转换完成之后会在刚才的目录下面生成一个**.engine**结尾的文件,如下

YOLOv7数据标注、环境配置、训练、预测、trt加速_第35张图片

这个就是我们最终需要的模型,将这个模型复制到yolov7里面和export.py同级的目录下

4.3、模型推理

yolov7目录下创建一个infer.py的文件,文件内容如下

import cv2
import tensorrt as trt
import torch
import numpy as np
from collections import OrderedDict, namedtuple


class TRT_engine():
    def __init__(self, weight) -> None:
        self.imgsz = [640, 640]
        self.weight = weight
        self.device = torch.device('cuda:0')
        self.init_engine()

    def init_engine(self):
        # Infer TensorRT Engine
        self.Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
        self.logger = trt.Logger(trt.Logger.INFO)
        trt.init_libnvinfer_plugins(self.logger, namespace="")
        with open(self.weight, 'rb') as self.f, trt.Runtime(self.logger) as self.runtime:
            self.model = self.runtime.deserialize_cuda_engine(self.f.read())
        self.bindings = OrderedDict()
        self.fp16 = False
        for index in range(self.model.num_bindings):
            self.name = self.model.get_binding_name(index)
            self.dtype = trt.nptype(self.model.get_binding_dtype(index))
            self.shape = tuple(self.model.get_binding_shape(index))
            self.data = torch.from_numpy(np.empty(self.shape, dtype=np.dtype(self.dtype))).to(self.device)
            self.bindings[self.name] = self.Binding(self.name, self.dtype, self.shape, self.data,
                                                    int(self.data.data_ptr()))
            if self.model.binding_is_input(index) and self.dtype == np.float16:
                self.fp16 = True
        self.binding_addrs = OrderedDict((n, d.ptr) for n, d in self.bindings.items())
        self.context = self.model.create_execution_context()

    def letterbox(self, im, color=(114, 114, 114), auto=False, scaleup=True, stride=32):
        # Resize and pad image while meeting stride-multiple constraints
        shape = im.shape[:2]  # current shape [height, width]
        new_shape = self.imgsz
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)
        # Scale ratio (new / old)
        self.r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        if not scaleup:  # only scale down, do not scale up (for better val mAP)
            self.r = min(self.r, 1.0)
        # Compute padding
        new_unpad = int(round(shape[1] * self.r)), int(round(shape[0] * self.r))
        self.dw, self.dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
        if auto:  # minimum rectangle
            self.dw, self.dh = np.mod(self.dw, stride), np.mod(self.dh, stride)  # wh padding
        self.dw /= 2  # divide padding into 2 sides
        self.dh /= 2
        if shape[::-1] != new_unpad:  # resize
            im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(self.dh - 0.1)), int(round(self.dh + 0.1))
        left, right = int(round(self.dw - 0.1)), int(round(self.dw + 0.1))
        self.img = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
        return self.img, self.r, self.dw, self.dh

    def preprocess(self, image):
        self.img, self.r, self.dw, self.dh = self.letterbox(image)
        self.img = self.img.transpose((2, 0, 1))
        self.img = np.expand_dims(self.img, 0)
        self.img = np.ascontiguousarray(self.img)
        self.img = torch.from_numpy(self.img).to(self.device)
        self.img = self.img.float()
        return self.img

    def predict(self, img, threshold):
        img = self.preprocess(img)
        self.binding_addrs['images'] = int(img.data_ptr())
        self.context.execute_v2(list(self.binding_addrs.values()))
        nums = self.bindings['num_dets'].data[0].tolist()
        boxes = self.bindings['det_boxes'].data[0].tolist()
        scores = self.bindings['det_scores'].data[0].tolist()
        classes = self.bindings['det_classes'].data[0].tolist()
        num = int(nums[0])
        new_bboxes = []
        for i in range(num):
            if (scores[i] < threshold):
                continue
            xmin = (boxes[i][0] - self.dw) / self.r
            ymin = (boxes[i][1] - self.dh) / self.r
            xmax = (boxes[i][2] - self.dw) / self.r
            ymax = (boxes[i][3] - self.dh) / self.r
            new_bboxes.append([classes[i], scores[i], xmin, ymin, xmax, ymax])
        return new_bboxes


def visualize(img, bbox_array):
    for temp in bbox_array:
        xmin = int(temp[2])
        ymin = int(temp[3])
        xmax = int(temp[4])
        ymax = int(temp[5])
        clas = int(temp[0])
        score = temp[1]
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (105, 237, 249), 2)
        img = cv2.putText(img, "class:" + str(clas) + " " + str(round(score, 2)), (xmin, int(ymin) - 5),
                          cv2.FONT_HERSHEY_SIMPLEX, 0.5, (105, 237, 249), 1)
    return img


trt_engine = TRT_engine("./best_fp16.engine")
# 使用摄像头进行实时预测
cap = cv2.VideoCapture(0)  # 0表示默认摄像头,如果有多个摄像头可以选择其他索引号
while True:
    ret, frame = cap.read()  # 读取一帧图像
    if not ret:
        break

    results = trt_engine.predict(frame, threshold=0.5)
    frame = visualize(frame, results)

    cv2.imshow("Camera", frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):  # 按下 'q' 键退出循环
        break

cap.release()  # 释放摄像头
cv2.destroyAllWindows()


# 预测单张图片
# img = cv2.imread("./data/coffee/images/1.jpg")
# results = trt_engine.predict(img, threshold=0.5)
# img = visualize(img, results)
# cv2.imshow("img", img)
# cv2.waitKey(0)

修改**.engine**路径之后直接运行即可,使用摄像头预测和单张图片预测只能开一个,另一个在运行的时候需要注释掉

你可能感兴趣的:(YOLO,python)