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

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

作者：Monodyee | 2024-03-23 14:02:57

赞

踩

yolov7训练

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)

image-20230720111441560

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

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

image-20230720112335609

1.2 下载预训练模型

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

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

image-20230720113021168

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

image-20230720120921032

2、环境准备

2.1 、python环境准备

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

2.1.1、使用anaconda安装python环境

下载地址 Free Download | Anaconda

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

单独配置环境变量

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

image-20230720155809951

添加如下环境变量

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

添加到用户path也可以

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

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

image-20230721150540387

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

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

image-20230720153545643

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

image-20230720153715677

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

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

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

配置环境变量

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

image-20230720155809951

添加如下环境变量

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

image-20230720160247242

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

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

image-20230720160529892

  1. cudnn

    下载地址

    cuDNN Download | NVIDIA Developer

image-20230720154133893

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

image-20230720160813139

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

2.2.2使用conda配置yolo环境

创建一个conda环境

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

image-20230727100823383

创建完成会有如下输出

image-20230727103423670

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

image-20230727103825320

切换到yolov7文件存放的位置

image-20230727104023451

注意里面的文件必须有 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()

  • 1
  • 2

image-20230727104817048

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

需要先卸载之后再安装

卸载指令：pip uninstall torch

进入pytorch官网，链接：PyTorchimage-20220515170214632

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

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

image-20220515170621528

到此，环境配置完成

2.3、使用cpu训练

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

2.3.1、使用conda配置yolo环境（cpu）

创建一个conda环境

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

image-20230727100823383

创建完成会有如下输出

image-20230727103423670

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

image-20230727103825320

切换到yolov7文件存放的位置

image-20230727104023451

注意里面的文件必须有 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运行即可

image-20230728094457454

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

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

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

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

image-20230728100950162

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

image-20230728102452222

3.2.1、划分数据集

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

使用下面代码进行划分,这个代码只划分了训练和验证，没有测试，注意修改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)

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38

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

image-20230728102214637

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']


  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9

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

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

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

修改train.py文件

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

image-20230801112023738

这些参数修改之后直接运行即可,显存小的或者使用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,修改以下参数

image-20230803095611759

只需修改对应的参数即可

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

注意版本

image-20230801143715376

4.1.1、linux配置

下载之后先解压

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

  • 1

配置环境变量

# 修改单个用户的环境变量文件
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

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8

然后安装对应的python包

image-20230801144840418

需要安装的python包分别在python、uuf、onnx_graphsurgeon里面，python里面需要安装的有三个，选择对应的python版本安装就行

安装指令

pip install *.whl

  • 1
4.1.2、windows配置

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

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

image-20230803102930878

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

接下来复制文件

  • 将 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包分别在python、uuf、onnx_graphsurgeon里面，python里面需要安装的有三个，选择对应的python版本安装就行

安装指令

pip install *.whl

  • 1

至此TensotRT已经配置完成。

4.2、模型转换

4.2.1、原生方法转engine
  1. .pt转**.onnx**

修改export.py，如下

image-20230803135131754

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

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

image-20230803135933462

在这个目录进入终端然后进入运行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**结尾的文件,如下

image-20230803141140456

这个就是我们最终需要的模型，将这个模型复制到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)

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115
  • 116
  • 117
  • 118
  • 119
  • 120
  • 121
  • 122
  • 123
  • 124
  • 125
  • 126
  • 127
  • 128
  • 129

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

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Monodyee/article/detail/295618
推荐阅读
相关标签

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

  

闽ICP备14008679号