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

YOLOv5 开源算法 -- 模型推理 detect.py，解析与修改_yolo开源组件

作者：weixin_40725706 | 2024-04-20 10:46:51

赞

踩

yolo开源组件

YOLOv5 开源算法 – 模型推理 detect.py，解析与修改

*** 源码有很多地方进行了修改，但大体运行逻辑不变
*** 本教程参考的教程链接：https://aitechtogether.com/python/79652.html
一、概述
  1. 目的： 去掉一些当前不常用的方法，将 detect.py 修改成方便python调用的模块，并精简为只支持图片/视频文件推理与。
二、代码解析
  1. 程序调用入口 main。
    def main(opt):
    run(**opt)

    • 1
    • 2
    1. 在其它地方运行YOLO时，只需要导入main即可,如下所示
    from yolov5.detect import main

    • 1

    1. 参数传递可以采用字典的格式传参，如下:

    2. 要注意字典的keys值必须与run方法想匹配；run函数中的参数方法下文有介绍。

    from yolov5.detect import main

config = {}
config["weights"] = "yolo5s.pt"

main(config)

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
  2. detect 中的模块包解析。
    1. python库模块与功能
      • python库一些常用功能模块。
        1. argparse： 它是一个用于命令项选项与参数解析的模块，通过在程序中定义好我们需要的参数，argparse 将会从 sys.argv 中解析出这些参数，并自动生成帮助和使用信息
        2. os： 它提供了多种操作系统的接口。通过os模块提供的操作系统接口，我们可以对操作系统里文件、终端、进程等进行操作
        3. sys： 它是与python解释器交互的一个接口，该模块提供对解释器使用或维护的一些变量的访问和获取，它提供了许多函数和变量来处理 Python 运行时环境的不同部分
        4. pathlib： 这个库提供了一种面向对象的方式来与文件系统交互，可以让代码更简洁、更易读
        5. torch： 这是主要的Pytorch库。它提供了构建、训练和评估神经网络的工具
        6. torch.backends. cudnn**：** 它提供了一个接口，用于使用cuDNN库，在NVIDIA GPU上高效地进行深度学习。cudnn模块是一个Pytorch库的扩展
      • import os           # 与操作系统进行交互的文件库 包含文件路径操作与解析
import platform     # 用于获取关于当前操作系统平台和硬件架构的信息
import sys          # sys模块包含了与python解释器和它的环境有关的函数。
from pathlib import Path    # Path能够更加方便得对字符串路径进行处理
import torch    			# pytorch 深度学习库
import shutil				# 文件处理，主要用于移动、复制文件

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
    2. 加载自定义模块
      • YOLO项目自定义的功能模块，主要作用有以下几个
        1. 将当前项目添加到系统路径上，以使得项目中的模块可以调用。
        2. 将当前项目的相对路径保存在ROOT中，便于寻找项目中的文件。
      • 部分模块的功能介绍如下：
        1. models.common.py： 这个文件定义了一些通用的函数和类，比如图像的处理、非极大值抑制等等。
        2. utils.dataloaders.py： 这个文件定义了两个类，LoadImages和LoadStreams，它们可以加载图像或视频帧，并对它们进行一些预处理，以便进行物体检测或识别。
        3. utils.general.py： 这个文件定义了一些常用的工具函数，比如检查文件是否存在、检查图像大小是否符合要求、打印命令行参数等等。
        4. utils.plots.py： 这个文件定义了Annotator类，可以在图像上绘制矩形框和标注信息。
        5. utils.torch_utils.py： 这个文件定义了一些与PyTorch有关的工具函数，比如选择设备、同步时间等等。
      • # __file__指的是当前文件路径(即detect.py),
# FILE最终保存着当前文件的绝对路径,比如D://yolov5/detect.py
FILE = Path(__file__).resolve() 
ROOT = FILE.parents[0]  # YOLOv5 root directory
# sys.path即当前python环境可以运行的路径,
# 假如当前项目不在该路径中,就无法运行其中的模块,
# 所以就需要加载路径
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
# relative ROOT设置为相对路径
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative  


''' 加载自定义模块 '''
from models.common import DetectMultiBackend
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, smart_inference_mode
import shutil

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
        • 7
        • 8
        • 9
        • 10
        • 11
        • 12
        • 13
        • 14
        • 15
        • 16
        • 17
        • 18
        • 19
        • 20
        • 21
    3. run 方法参数解析
      • run：配置运行参数，调用模型开始推理文件；这些参数的功能介绍如下：
        1. weights： 训练的权重路径，可以使用自己训练的权重，也可以使用官网提供的权重。默认官网的权重yolov5s.pt(yolov5n.pt/yolov5s.pt/yolov5m.pt/yolov5l.pt/yolov5x.pt/区别在于网络的宽度和深度以此增加)
        2. source： 测试数据，可以是图片/视频路径，也可以是’0′(电脑自带摄像头)，也可以是rtsp等视频流, 默认data/images
        3. data： 配置数据文件路径，包括image/label/classes等信息，训练自己的文件，需要作相应更改，可以不用管
        4. imgsz： 预测时网络输入图片的尺寸，默认值为 [640]
        5. conf_thres： 置信度阈值，默认为 0.50
        6. iou_thres： 非极大抑制时的 IoU 阈值，默认为 0.45
        7. max_det： 保留的最大检测框数量，每张图片中检测目标的个数最多为1000类
        8. device： 使用的设备，可以是 cuda 设备的 ID（例如 0、0,1,2,3）或者是 ‘cpu’，默认为 ‘0’
        9. view_img： 是否展示预测之后的图片/视频，默认False
        10. save_txt： 是否将预测的框坐标以txt文件形式保存，默认False，使用–save-txt 在路径runs/detect/exp*/labels/*.txt下生成每张图片预测的txt文件
        11. save_conf： 是否保存检测结果的置信度到 txt文件，默认为 False
        12. save_crop： 是否保存裁剪预测框图片，默认为False，使用–save-crop 在runs/detect/exp*/crop/剪切类别文件夹/ 路径下会保存每个接下来的目标
        13. nosave： 不保存图片、视频，要保存图片，不设置–nosave 在runs/detect/exp*/会出现预测的结果
        14. classes： 仅检测指定类别，默认为 None
        15. agnostic_nms： 是否使用类别不敏感的非极大抑制（即不考虑类别信息），默认为 False
        16. augment： 是否使用数据增强进行推理，默认为 False
        17. visualize： 是否可视化特征图，默认为 False
        18. update： 如果为True，则对所有模型进行strip_optimizer操作，去除pt文件中的优化器等信息，默认为False
        19. project： 结果保存的项目目录路径，默认为 ‘ROOT/runs/detect’
        20. name： 结果保存的子目录名称，默认为 ‘exp’
        21. exist_ok： 是否覆盖已有结果，默认为 False
        22. line-thickness： 画 bounding box 时的线条宽度，默认为 3
        23. hide_labels： 是否隐藏标签信息，默认为 False
        24. hide_conf： 是否隐藏置信度信息，默认为 False
        25. half： 是否使用 FP16 半精度进行推理，默认为 False
        26. dnn： 是否使用 OpenCV DNN 进行 ONNX 推理，默认为 False
        27. vid_stride: video frame-rate stride
        28. class_dir: 指定被识别的原图保存路径
        29. save_crop_path: 指定新的截图保存路径
        30. **save_label_path:**指定标签保存路径
        31. iou_configs:辅助检测结果
      • run(
        # model.pt path(s) 事先训练完成的权重文件，比如yolov5s.pt,默认 weights/，
        # 假如使用官方训练好的文件（比如yolov5s）,则会自动下载
        weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
        # file/dir/URL/glob, 0 for webcam 预测时的输入数据，可以是文件/路径/URL/glob, 
        # 输入是0的话调用摄像头作为输入，默认data/images/
        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam
        # dataset.yaml path, data文件路径，包括类别/图片/标签等信息
        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path
        # inference size (pixels)  
        # 预测时的放缩后图片大小(因为YOLO算法需要预先放缩图片), 
        # 两个值分别是height, width。默认640*640
        imgsz=(640, 640),  # inference size (height, width)
        # confidence threshold 置信度阈值, 高于此值的bounding_box才会被保留。
        # 默认0.25，用在nms中
        conf_thres=0.80,  # confidence threshold
        # NMS IOU threshold IOU阈值
        # 高于此值的bounding_box才会被保留。默认0.45，用在nms中
        iou_thres=0.45,  # NMS IOU threshold
        # maximum detections per image 
        # 一张图片上检测的最大目标数量，用在nms中
        max_det=1000,  # maximum detections per image
        # cuda device, i.e. 0 or 0,1,2,3 or cpu 所使用的GPU编号，
        # 如果使用CPU就写cpu
        device='0',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        # show results 是否展示预测之后的图片或视频，
        # 默认False
        view_img=False,  # show results
        # save results to *.txt 是否将预测的框坐标以txt文件形式保存, 默认False, 
        # 使用--save-txt 在路径runs/detect/exp*/labels/*.txt下生成每张图片预测的txt文件
        save_txt=True,  # save results to *.txt
        # save confidences in --save-txt labels 
        # 是否将结果中的置信度保存在txt文件中，默认False
        save_conf=False,  # save confidences in --save-txt labels
        # save cropped prediction boxes 是否保存裁剪后的预测框，默认为False, 
        # 使用--save-crop 在runs/detect/exp*/crop/剪切类别文件夹/ 路径下会保存每个接下来的目标
        save_crop=True,  # save cropped prediction boxes
        # do not save images/videos 不保存图片、视频, 要保存图片，
        # 不设置--nosave 在runs/detect/exp*/会出现预测的结果
        nosave=False,  # do not save images/videos
        # filter by class: --class 0, or --class 0 2 3 过滤指定类的预测结果
        classes=None,  # filter by class: --class 0, or --class 0 2 3
        # class-agnostic NMS 进行NMS去除不同类别之间的框, 默认False
        agnostic_nms=False,  # class-agnostic NMS
        # augmented inference TTA测试时增强/多尺度预测，可以提分
        augment=False,  # augmented inference
        # visualize features 是否可视化网络层输出特征
        visualize=False,  # visualize features
        # update all models 如果为True,则对所有模型进行strip_optimizer操作,去除pt文件中的优化器等信息,
        # 默认为False
        update=False,  # update all models
        # save results to project/name 预测结果保存的路径
        project=ROOT / 'runs/detect',  # save results to project/name
        # save results to project/name 结果保存文件夹的命名前缀
        name='exp',  # save results to project/name
        # existing project/name ok, do not increment True: 推理结果覆盖之前的结果 
        # False: 推理结果新建文件夹保存,文件夹名递增
        exist_ok=False,  # existing project/name ok, do not increment
        # bounding box thickness (pixels) 绘制Bounding_box的线宽度
        line_thickness=3,  # bounding box thickness (pixels)
        # hide labels 若为True: 隐藏标签
        hide_labels=False,  # hide labels
        # hide confidences  若为True: 隐藏置信度
        hide_conf=False,  # hide confidences
        # use FP16 half-precision inference 是否使用半精度推理（节约显存）
        half=False,  # use FP16 half-precision inference
        # use OpenCV DNN for ONNX inference 是否使用OpenCV DNN预测
        dnn=False,  # use OpenCV DNN for ONNX inference
        
        vid_stride=1,  # video frame-rate stride
        class_dir="", 	# 指定被识别的原图保存路径
        save_crop_path = None,	# 指定新的截图保存路径
        save_label_path = None,	# 指定标签保存路径
        iou_configs = None		# iou 辅助检测结果

):

        • 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
    4. 初始化配置
      1. 主要用于处理输入来源。定义了一些布尔值区分输入是图片、视频.
        1. save_img： 判断是否保存图片和txt文件，如果nosave(传入的参数)为false且source的结尾不是txt则保存图片
        2. **is_file：**判断source是不是视频/图像文件路径
          1. Path() 提取文件名，suffix：最后一个组件的文件扩展名。例：若source是"D://YOLOv5/data/1.jpg"， 则Path(source).suffix是".jpg"， Path(source).suffix[1:]是"jpg"
          2. IMG_FORMATS 和 VID_FORMATS两个变量保存的是所有的视频和图片的格式后缀。
      2. source = str(source) # 输入的路径变为字符串	

# 是否保存图片和txt文件，如果nosave(传入的参数)为false且source的结尾不是txt则保存图片
save_img = not nosave and not source.endswith('.txt')  # save 
# 判断source是不是视频/图像文件路径
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
    5. 保存结果
      1. 主要是用于创建保存输出结果的目录。创建一个新的文件夹exp（在runs文件夹下）来保存运行的结果。
        1. project 指 run 函数中的 project，对应的是 runs/detect 的目录，name 对应 run 函数中的“name=exp”，然后进行拼接操作。
        2. 使用increment_path函数来确保目录不存在，如果存在，则在名称后面添加递增的数字。
        3. 然后判断 save_txt 是否为 true，save_txt 在 run 函数以及 parse_opt() 函数中都有相应操作，如果传入save_txt，新建 “labels” 文件夹存储结果.
      2. # save_dir是保存运行结果的文件夹名，
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) 
# 如果目录已经存在，而exist_ok为False**，那么会抛出一个异常，指示目录已存在。
# 如果exist_ok为True，则不会抛出异常，而是直接使用已经存在的目录
(save_label_path if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
save_dir.mkdir(parents=True, exist_ok=True)  # make dir

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
    6. 加载模型

      1. 主要是用于选择设备、初始化模型和检查图像大小。

        1. 首先调用select_device函数选择设备，如果device为空，则使用默认设备。

        2. 然后使用DetectMultiBackend类来初始化模型，其中

          • weights 指模型的权重路径
          • device 指设备
          • dnn 指是否使用OpenCV DNN
          • data 指数据集配置文件的路径
          • fp16 指是否使用半精度浮点数进行推理

          接着从模型中获取stride、names和pt等参数，其中

          • stride 指下采样率
          • names 指模型预测的类别名称
          • pt 是Pytorch模型对象

        3. 最后调用check_img_size函数检查图像大小是否符合要求，如果不符合则进行调整。

      2. # 获取设备 CPU/CUDA
device = select_device(device)
# DetectMultiBackend定义在models.common模块中，是我们要加载的网络，
# 其中weights参数就是输入时指定的权重文件（比如yolov5s.pt）
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
'''
    stride：推理时所用到的步长，默认为32， 大步长适合于大目标，小步长适合于小目标
    names：保存推理结果名的列表，比如默认模型的值是['person', 'bicycle', 'car', ...]
    pt: 加载的是否是pytorch模型（也就是pt格式的文件）
    jit：当某段代码即将第一次被执行时进行编译，因而叫“即时编译”
    onnx：利用Pytorch我们可以将model.pt转化为model.onnx格式的权重，在这里onnx充当一个后缀名称，
          model.onnx就代表ONNX格式的权重文件，这个权重文件不仅包含了权重值，也包含了神经网络的网络流动信息以及每一层网络的输入输出信息和一些其他的辅助信息。
'''
stride, names, pt = model.stride, model.names, model.pt
# 确保输入图片的尺寸imgsz能整除stride=32 如果不能则调整为能被整除并返回
imgsz = check_img_size(imgsz, s=stride)  # check image size

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
        • 7
        • 8
        • 9
        • 10
        • 11
        • 12
        • 13
        • 14
        • 15
        • 16
    7. 加载数据

      1. 直接从source文件下读取图片使用 LoadImages 加载图像

      2. # 直接从source文件下读取图片
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
bs = 1  # batch_size
# 保存视频的路径
# 前者是视频路径,后者是一个cv2.VideoWriter对象
vid_path, vid_writer = [None] * bs, [None] * bs

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
    8. 开始推理

      1. 推理部分是整个算法的核心部分。通过for循环对加载的数据进行遍历，一帧一帧地推理，进行NMS非极大值抑制、绘制bounding box、预测类别。

        1. 热身部分

          1. 这段代码进行了模型的热身（warmup）操作，即对模型进行一些预处理以加速后续的推理过程。

            1. 首先定义了一些变量，包括seen、windows和dt，分别表示已处理的图片数量、窗口列表和时间消耗列表。遍历 dataset ，整理图片信息；进行预测，根据 run 函数里面的置信度以及IOU参数，进行信息过滤；对检测框进行后续处理，画框选择，坐标映射（640*640坐标映射为原图坐标），是否保存绘画结果。
            2. 接着对数据集中的每张图片进行预处理：
              • 首先将图片转换为Tensor格式，并根据需要将其转换为FP16或FP32格式。
              • 然后将像素值从0-255转换为0.0-1.0归一化，并为批处理增加一维。
              • 最后记录时间消耗并更新dt列表。
          2. # Run inference
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz))  # warmup
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())

# 去遍历图片，进行计数，
for path, im, im0s, vid_cap, s in dataset:
    '''
     在dataset中，每次迭代的返回值是self.sources, img, img0, None, ''
      path：文件路径（即source）
      im: resize后的图片（经过了放缩操作）
      im0s: 原始图片
      vid_cap=none
      s： 图片的基本信息，比如路径，大小
    '''
    with dt[0]:
        # ===以下部分是做预处理===#

        # 将图片放到指定设备(如GPU)上识别。#torch.size=[3,640,480]
        im = torch.from_numpy(im).to(device)   
        # uint8 to fp16/32 # 把输入从整型转化为半精度/全精度浮点数。 
        im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
        # 0 - 255 to 0.0 - 1.0 归一化，所有像素点除以255
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            # expand for batch dim 添加一个第0维。
            # 缺少batch这个尺寸，所以将它扩充一下，变成[1，3,640,480]
            im = im[None]  # expand for batch dim

            • 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
        2. 对每张图片/视频进行前向推理

          1. 这段代码对每张图片/视频进行前向推理

            1. 第一行代码，创建了一个名为**“visualize”的变量，如果需要可视化，则将其设置为保存可视化结果的路径，否则将其设置为False。使用increment_path函数**创建路径，如果文件名已存在，则将数字附加到文件名后面以避免覆盖已有文件。

            2. 第二行代码，使用model函数对图像im进行预测，augment和visualize参数用于指示是否应该在预测时使用数据增强和可视化。

            3. 第三行代码，记录了当前时间，并计算从上一个时间点到这个时间点的时间差，然后将这个时间差加到一个名为dt的时间差列表中的第二个元素上。

            4. # Inference
with dt[1]:
    # 可视化文件路径。如果为True则保留推理过程中的特征图，保存在runs文件夹中
    visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
    # 推理结果，pred保存的是所有的bound_box的信息，
    # 模型预测出来的所有检测框，torch.size=[1,18900,85]
    pred = model(im, augment=augment, visualize=visualize)

              • 1
              • 2
              • 3
              • 4
              • 5
              • 6
              • 7
        3. NMS除去多余的框

          1. 这段代码是执行**非最大值抑制（NMS）**的步骤，用于筛选预测结

            1. non_max_suppression函数的输入参数包括预测结果pred、置信度阈值conf_thres、IOU（交并比）阈值iou_thres、类别classes、是否进行类别无关的NMSagnostic_nms，以及最大检测数max_det。该函数的输出是经过NMS筛选后的预测结果。
            2. 第二行代码更新了计时器，记录了NMS操作所用的时间。
          2. # NMS
with dt[2]:
    # 执行非极大值抑制，返回值为过滤后的预测框
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
    '''
        pred: 网络的输出结果
        conf_thres： 置信度阈值
        iou_thres： iou阈值
        classes: 是否只保留特定的类别 默认为None
        agnostic_nms： 进行nms是否也去除不同类别之间的框
        max_det: 检测框结果的最大数量 默认1000
    '''
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
            • 10
            • 11
            • 12
            • 13
            • 14
        4. 预测过程

          1. 这段代码使用了一个循环来遍历检测结果列表中的每个物体，并对每个物体进行处理。

          2. 循环中的变量**“i”是一个索引变量，表示当前正在处理第几个物体，而变量“det”则表示当前物体的检测结果**。循环体中的第一行代码 “seen += 1” 用于增加一个计数器，记录已处理的物体数量。

          3. 接下来，根据是否使用网络摄像头来判断处理单张图像还是批量图像。

          4. *如果使用的是网络摄像头，则代码会遍历每个图像并复制一份备份到变量”im0″中，同时将当前图像的路径和计数器记录到变量“p”和“frame”中。最后，将当前处理的物体索引和相关信息记录到字符串变量”s”中。

          5. 如果没有使用网络摄像头，**则会直接使用“im0s”变量中的图像，将图像路径和计数器记录到变量“p”和“frame”中。**同时，还会检查数据集中是否有“frame”属性，如果有，则将其值记录到变量“frame”中。

            iou = 0
# Process predictions
# 把所有的检测框画到原图中
'''
   i：每个batch的信息
   det:表示5个检测框的信息
'''
for i, det in enumerate(pred):  # per image
    seen += 1   # seen是一个计数的功能
    '''
   从LoadImages流读取本都文件中的照片或者视频 所以batch_size=1
      p: 当前图片/视频的绝对路径 如 F:\yolo_v5\yolov5-U\data\images\bus.jpg
      s: 输出信息 初始为 ''
      im0: 原始图片 letterbox + pad 之前的图片
      frame: 视频流,此次取的是第几张图片
    '''
    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)


            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
            • 10
            • 11
            • 12
            • 13
            • 14
            • 15
            • 16
            • 17
            • 18

          6. 将图像路径转换为**“Path”**对象。

            1. 接下来，**使用“save_dir”变量中的路径和图像文件名来构建保存检测结果图像的完整路径，并将其保存在变量“save_path”中。**根据数据集的模式（”image”或”video”）来构建保存检测结果标签的文件路径，并将其保存在变量“txt_path”中。
            2. 在处理图像路径和文件路径之后，将图像的尺寸信息添加到字符串变量“s”中，以便于打印。接着，计算归一化增益“gn”，并将其保存在变量中，以便后续使用。
            3. 根据是否需要保存截取图像的标志“save_crop”来选择是否要对原始图像进行复制**，以备保存截取图像时使用。**
            4. **最后，创建了一个“Annotator”对象，以便于在图像上绘制检测结果。
            p = Path(p)  # to Path
# 图片/视频的保存路径save_path 如 runs\\detect\\exp8\\fire.jpg
save_path = str(save_dir / p.name)  # im.jpg
# 设置保存框坐标的txt文件路径，每张图片对应一个框坐标信息
txt_path = str(Path(save_label_path) / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
# 设置输出图片信息。图片shape (w, h)
s += '%gx%g ' % im.shape[2:]  # print string
# 得到原图的宽和高
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
# 保存截图。如果save_crop的值为true，则将检测到的bounding_box单独保存成一张图片。
imc = im0.copy() if save_crop else im0  # for save_crop
# 得到一个绘图的类，类中预先存储了原图、线条宽度、类名
annotator = Annotator(im0, line_width=line_thickness, example=str(names))

            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
            • 10
            • 11
            • 12
            • 13

          7. 这段代码会判断有没有框。如果检测结果列表中存在物体，则代码会执行一些操作。

            1. 首先，将检测结果中的物体坐标从缩放后的图像大小还原回原始图像的大小。这里使用了一个名为“scale_coords”的函数来进行缩放，该函数的作用是将物体坐标从缩放前的大小变换到缩放后的大小
            2. 接着，遍历每个物体，**将其类别和数量添加到字符串变量“s”中。**具体来说，计算当前类别下检测到的物体数量“n”，然后根据数量和类别名字构建一段字符串，并将其添加到变量“s”中。代码中的“names”变量包含了数据集中所有类别的名称。
            3. 最后，返回字符串变量“s”，并结束当前代码块。
            if len(det):    # 判断有没有框
    # Rescale boxes from img_size to im0 size
    # Rescale boxes from img_size to im0 size
    # 将预测信息映射到原图
    # 将标注的bounding_box大小调整为和原图一致（因为训练时原图经过了放缩）此时坐标格式为xyxy
    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

    # Print results
    # 打印检测到的类别数量
    for c in det[:, 5].unique():
        n = (det[:, 5] == c).sum()  # detections per class
        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
        # LOGGER.info(f'{s}Done. {dt[1].dt * 1E3:.1f}ms')
        LOGGER.info(f"{s.split(' ')[0]} {s.split(' ')[1]} {n} {names[int(c)]}{'s' * (n > 1)} {dt[1].dt * 1E3:.1f}ms")
                    

            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
            • 10
            • 11
            • 12
            • 13
            • 14
            • 15
        5. 打印目标检测结果

          1. 这段代码是打印目标检测结果的一些操作

            1. 将存在检测结果的原图移动到指定路径
            2. 如果需要将检测结果写入文件，则将检测结果中的物体坐标转换为相对于原始图像的归一化坐标，并将其写入到以图像文件名命名的“.txt”文件中
            3. 如果需要保存检测结果图像或者在图像上绘制框**，**每个物体添加一个边界框，并将其标记在图像上
            4. 如果需要将边界框截取出来保存至指定的目录下保存*
          2. #  将原图移动过去
if class_dir:
    dest_dir = os.path.join(class_dir, str(names[int(c)]))
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)

    shutil.copy(p, dest_dir)

# 保存预测结果：txt/图片画框/crop-image
for *xyxy, conf, cls in reversed(det):
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 
# 每行: class_id + score + xywh
if save_txt:  # Write to file   # Write to file 保存txt文件
# 将xyxy(左上角+右下角)格式转为xywh(中心点+宽长)格式，并归一化，转化为列表再保存
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
# line的形式是： ”类别 x y w h“，若save_conf为true，则line的形式是：”类别 x y w h 置信度“
line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format

with open(txt_path + '.txt', 'a') as f:
    # 写入对应的文件夹里，路径默认为 “runs\detect\exp*\labels”
    f.write(('%g ' * len(line)).rstrip() % line + '\n')

# 在原图上画框+将预测到的目标剪切出来保存成图片，保存在save_dir/crops下，在原图像画图或者保存结果
if save_img or save_crop or view_img:  # Add bbox to image
# 类别标号
c = int(cls)  # integer class   #
# 类别名
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

#绘制边框
annotator.box_label(xyxy, label, color=colors(c, True))

# 在原图上画框+将预测到的目标剪切出来保存成图片，保存在save_dir/crops下（单独保存）
if save_crop and save_crop_path:
    save_one_box(xyxy, imc, file= Path(save_crop_path) / names[c] / f'{p.stem}.jpg', BGR=True)

            • 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
        6. 在窗口中实时查看检测结果

          1. 这段代码是实现在输出窗口实时查看检测结果

            如果需要在窗口中实时查看检测结果，则会使用OpenCV库中的函数将图像显示在窗口中，并等待1毫秒以便继续下一帧的检测。

            **代码会检查是否已经为当前图像创建了窗口（if p not in windows），并在必要时创建窗口，并使用图像名称来命名该窗口。**窗口的名称是由变量“p”指定的图像路径名。

            • 如果检测到图像尚未在窗口中打开，则代码会创建一个新窗口并将图像显示在窗口中。
            • 如果图像已经在窗口中打开，则代码会直接更新窗口中的图像。
          2. # 如果设置展示，则show图片 / 视频
im0 = annotator.result()    # im0是绘制好的图片
if view_img:    # 显示图片
    if platform.system() == 'Linux' and p not in windows:
        windows.append(p)
        cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
        cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
    cv2.imshow(str(p), im0)
    cv2.waitKey(1)  # 1 millisecond

            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
        7. 设置保存结果

          1. 这段代码是设置保存图片和视频

            首先“save_img”判断是否是图片，如果是则保存路径和图片；如果是视频或流，需要重新创建视频文件。

            保存视频文件中，

            • 如果是视频“vid_cap”，则使用python-opencv读取视频，计算视频帧速率FPS以及视频帧宽度和高度。
            • 如果是流，则保存路径后缀加上’.mp4′
          2. # 设置保存图片/视频
if save_img:    # 如果save_img为true,则保存绘制完的图片
    if dataset.mode == 'image':  # 如果是图片,则保存
        cv2.imwrite(save_path, im0)
        # cv2.imwrite(save_path, imc)
    else:  # 'video' or 'stream'    # 'video' or 'stream'  如果是视频或者"流"
        if vid_path[i] != save_path:  # new video
            vid_path[i] = save_path
            # 以下的部分是保存视频文件
            if isinstance(vid_writer[i], cv2.VideoWriter):
                vid_writer[i].release()  # release previous video writer
            if vid_cap:  # video
                fps = vid_cap.get(cv2.CAP_PROP_FPS)
                w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            else:  # stream
                fps, w, h = 30, im0.shape[1], im0.shape[0]
            save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
            vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        vid_writer[i].write(im0)

            • 1
            • 2
            • 3
            • 4
            • 5
            • 6
            • 7
            • 8
            • 9
            • 10
            • 11
            • 12
            • 13
            • 14
            • 15
            • 16
            • 17
            • 18
            • 19
            • 20
        8. 来对推理过程这一部分代码做一个总结：

          这一段代码是一个目标检测算法中的推理过程，通过对一张或多张图片中的物体进行检测，输出检测结果，并将检测结果保存到文件或显示在窗口中。以下是每个步骤的详细说明：

          1. 对于每个输入图片，将其路径、原始图像和当前帧数（如果存在）分别赋值给p、im0和frame变量；
          2. 如果webcam为True，则将输出信息字符串s初始化为空，否则将其初始化为该数据集的“frame”属性；
          3. 将p转换为Path类型，并生成保存检测结果的路径save_path和文本文件路径txt_path；
          4. 将im0大小与目标检测的输入大小匹配，将检测结果det中的边界框坐标从img_size缩放到im0大小，然后将结果打印在输出字符串s中；
          5. 如果save_txt为True，则将结果写入文本文件中；
          6. 如果save_img、save_crop或view_img中任意一个为True，则将检测结果添加到图像中，并在窗口中显示结果；
          7. 如果save_img为True，则保存结果图像；
          8. 如果是视频数据集，则将结果写入视频文件中；
          9. 最后，打印每个图片的检测时间。
    9. 在终端里打印出运行的结果

      1. 这部分代码用于打印结果，记录了一些总共的耗时，以及信息保存。

        输出结果包括每张图片的预处理、推理和NMS时间，以及结果保存的路径。

        如果update为True，则将模型更新，以修复SourceChangeWarning。

      2. # 平均每张图片所耗费时间
t = tuple(x.t / seen * 1E3 for x in dt)  # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_img:
    s = f"\n{len(list(save_label_path.glob('*.txt')))} labels saved to {save_label_path}" if save_txt else ''
    LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
    strip_optimizer(weights[0])  # update model (to fix SourceChangeWarning)

        • 1
        • 2
        • 3
        • 4
        • 5
        • 6
        • 7
        • 8
三、完整源码
# YOLOv5 
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/weixin_40725706/article/detail/456911
推荐阅读
相关标签

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