detect.py
detect.py
目录
1.所需的库和模块
import argparse
import time
from pathlib import Path
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
from utils.plots import plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
2.def detect(save_img=False):
# 这段代码是一个完整的 detect 函数,它实现了使用深度学习模型进行图像或视频的检测任务。
# 函数定义。这个函数名为 detect ,接受一个参数。
# 1.save_img :默认值为 False ,表示是否保存推理后的图像。
def detect(save_img=False):
# 参数提取。这里从 opt 对象中提取了几个参数 :
# source :输入源,可以是图像文件路径、视频文件路径或摄像头索引。
# weights :模型权重文件的路径。
# view_img :是否在推理过程中查看图像。
# save_txt :是否保存包含检测结果的文本文件。
# imgsz :输入图像的大小。
# trace :是否启用模型的追踪功能。
source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, opt.trace
# 保存图像逻辑。这行代码重新定义了 save_img 的值。如果 opt.nosave 为 False (即允许保存图像)且 source 不以 .txt 结尾(即不是文本文件),则 save_img 被设置为 True ,表示会保存推理后的图像。
save_img = not opt.nosave and not source.endswith('.txt') # save inference images 保存推理图像。
# 网络摄像头或视频流检测。这行代码用于判断输入源是否为网络摄像头或视频流。
# 如果 source 是数字(代表摄像头索引)、以 .txt 结尾(代表视频流列表文件),或者以 rtsp:// 、 rtmp:// 、 http:// 或 https:// 开头的字符串(代表视频流URL),则 webcam 被设置为 True 。
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
# 这段代码是检测函数的初始化部分,它设置了检测任务的基本参数和条件,为后续的图像处理和模型推理做准备。
# Directories
# 设置保存目录。
# 这段代码是 detect 函数中负责设置保存目录的部分。它使用 Python 的 pathlib 库来处理文件路径,并确保所需的目录存在。
# Path(opt.project) / opt.name :这里使用 pathlib 的 Path 类来构建一个路径。 opt.project 指定了项目的根目录, opt.name 指定了项目中的一个特定名称,两者结合形成了一个具体的路径。
# increment_path() :这是一个自定义函数,它的作用是为给定的路径生成一个唯一的版本。如果路径已存在,它会增加一个数字后缀,以确保路径的唯一性。
# exist_ok=opt.exist_ok 参数控制如果路径已存在时的行为 :如果 opt.exist_ok 为 True ,则不增加后缀,允许覆盖;如果为 False ,则增加后缀以避免覆盖。
# def increment_path(path, exist_ok=True, sep=''):
# -> 用于生成一个唯一的文件路径。如果指定的路径已经存在,函数会通过添加一个数字后缀来创建一个新的路径。返回一个新的路径,它是原始路径加上一个更新的数字后缀。
# -> return f"{path}{sep}{n}" # update path
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
# 创建目录。
# (save_dir / 'labels' if save_txt else save_dir) :这里使用条件表达式来决定最终的目录路径。如果 save_txt 为 True (即需要保存文本标签),则在 save_dir 下创建一个名为 labels 的子目录;否则,直接使用 save_dir 。
# mkdir(parents=True, exist_ok=True) :使用 mkdir 方法创建目录。 parents=True 表示如果父目录不存在,则一并创建; exist_ok=True 表示如果目录已存在,则不抛出异常。
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# 这段代码的主要功能是确保检测结果的保存目录存在,并且如果需要保存文本标签,则创建一个额外的 labels 子目录。通过这种方式,检测结果(如图像和文本文件)可以被组织得井井有条,便于后续的访问和管理。
# Initialize
# 这段代码是 detect 函数中用于初始化的部分,它负责设置日志记录和选择计算设备,并且还检查是否使用半精度计算。
# 日志记录设置。这行代码调用了一个名为 set_logging 的函数,这个函数用来配置日志记录器的行为,比如设置日志级别、日志格式和日志输出的位置。
# def set_logging(rank=-1): -> 用于配置 Python 的日志记录系统。函数的目的是在程序开始时设置合适的日志记录级别。在分布式训练环境中, rank 参数可以用来区分主进程和其他进程,以便在非主进程中减少日志输出,避免日志信息的冗余。
set_logging()
# 设备选择。 select_device 是一个自定义函数,用于选择用于模型推理的设备。这个函数根据 opt.device 参数来决定是使用 CPU 还是 GPU。 opt.device 是一个字符串,如 'cpu' 或 'cuda:0' (表示第一个 GPU),或者其他 GPU 设备标识。
# def select_device(device='', batch_size=None): -> 用于选择并配置 PyTorch 模型将使用的计算设备,可以是 CPU 或者一个或多个 GPU。返回一个 torch.device 对象,表示选择的设备。 -> return torch.device('cuda:0' if cuda else 'cpu')
device = select_device(opt.device)
# 半精度计算检查。这行代码检查选择的设备是否为 CPU。如果 device.type 不等于 'cpu' ,即表示使用的是 CUDA 设备(GPU),那么 half 变量被设置为 True ,表示可以启用半精度(FP16)计算。半精度计算只在 CUDA 设备上支持,可以加速模型推理并减少内存使用,但可能牺牲一些精度。
half = device.type != 'cpu' # half precision only supported on CUDA 半精度仅支持 CUDA。
# 这段代码的目的是为模型推理准备合适的环境 :
# 通过 set_logging 配置日志记录,以便在推理过程中记录重要信息。
# 通过 select_device 选择一个合适的计算设备,根据设备类型决定是否可以使用 GPU 加速。
# 根据选择的设备决定是否启用半精度计算,以优化性能和资源使用。
# 这些初始化步骤是进行深度学习模型推理前的重要准备工作,确保模型可以在正确的设备上以合适的精度运行。
# Load model
# 这段代码是 detect 函数中用于加载和准备模型的部分。
# 模型加载。
# attempt_load 是一个自定义函数,用于尝试加载模型权重。它接受两个参数: weights (模型权重文件的路径)和 map_location (指定模型加载到的设备,如 CPU 或 GPU)。
# 这行代码加载模型,并默认为 FP32(单精度浮点数)模型。
# def attempt_load(weights, map_location=None): -> 负责加载一个或多个模型权重,并将它们组合成一个集成模型(Ensemble)。返回包含多个模型的模型集成对象。 -> return model # return ensemble
model = attempt_load(weights, map_location=device) # load FP32 model
# 获取模型步长。这行代码获取模型的步长(stride),即模型在输入图像上滑动窗口的步长。 model.stride 返回一个包含步长的张量, max() 函数找到最大步长, int() 确保步长为整数。
stride = int(model.stride.max()) # model stride
# 检查图像尺寸。check_img_size 是一个自定义函数,用于确保提供的图像尺寸 imgsz 与模型的步长 stride 兼容。如果需要,这个函数可能会调整图像尺寸以适应模型。
# def check_img_size(img_size, s=32):
# -> 目的是验证给定的图像尺寸 img_size 是否是步长 s 的倍数。如果不是,函数将调整 img_size 到最接近的、大于或等于 img_size 的 s 的倍数,并打印一条警告信息。返回调整后的图像尺寸 new_size 。
# -> return new_size
imgsz = check_img_size(imgsz, s=stride) # check img_size
# 模型追踪。
if trace:
# 如果 trace 为 True ,则使用 TracedModel 包装原始模型。 TracedModel 是一个自定义类,用于追踪模型的执行,通常用于性能优化和分析。这里它接受原始模型、设备和图像尺寸作为参数。
# class TracedModel(nn.Module):
# -> TracedModel 的目的是将一个给定的 PyTorch 模型转换为一个追踪(traced)模型,也称为 TorchScript 模型。
# -> def __init__(self, model=None, device=None, img_size=(640,640)):
# -> def forward(self, x, augment=False, profile=False):
model = TracedModel(model, device, opt.img_size)
# 转换为半精度。
if half:
# 如果 half 为 True ,即设备支持半精度计算(通常是 GPU),则调用 model.half() 将模型转换为 FP16(半精度浮点数)。这可以减少模型的内存占用并可能加速推理过程,但可能会牺牲一些精度。
model.half() # to FP16
# 这段代码的主要功能是加载模型,确保模型与输入图像尺寸兼容,并根据设备能力对模型进行优化(如转换为半精度)。这些步骤是模型推理前的重要准备,确保模型可以在正确的设备上以最佳性能运行。
# Second-stage classifier
# 这段代码是用于在目标检测流程中添加第二阶段分类器的示例。第二阶段分类器通常用于对第一阶段检测器识别的目标进行更精确的分类。
# 分类器启用标志。这里定义了一个变量 classify 并将其设置为 False ,这意味着默认情况下不会启用第二阶段分类器。如果需要启用分类器,你需要将这个变量设置为 True 。
classify = False
# 分类器初始化和加载。
if classify:
# if classify : 这个条件判断基于 classify 变量的值。由于 classify 默认为 False ,所以下面的代码块在默认情况下不会被执行。
# modelc = load_classifier(name='resnet101', n=2) :如果 classify 为 True ,这行代码将调用 load_classifier 函数来初始化一个名为 'resnet101' 的分类器模型, n=2 表示模型的输出类别数为 2。
# def load_classifier(name='resnet101', n=2): -> 它用于加载一个预训练的图像分类模型,并将其输出层重塑为 n 类输出。返回一个重塑了输出层的预训练模型。 -> return model
modelc = load_classifier(name='resnet101', n=2) # initialize
# modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) :这行代码使用 PyTorch 的 torch.load 函数加载预训练的 ResNet-101 模型权重。
# 'weights/resnet101.pt' 是权重文件的路径, map_location=device 指定了模型权重加载到的设备(CPU 或 GPU)。
# .to(device) :将分类器模型移动到之前选择的设备(CPU 或 GPU)上。 .eval() :将模型设置为评估模式,这是在模型推理时必须的,以确保模型的行为符合推理模式,例如关闭 dropout 等仅在训练时使用的层。
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
# 这段代码提供了一个可选的第二阶段分类器的初始化和加载过程。如果你的目标检测任务需要在检测到目标后进行更精确的分类,你可以将 classify 变量设置为 True ,并确保 load_classifier 函数和权重文件路径正确无误。
# 这样,你就可以在目标检测流程中集成一个预训练的 ResNet-101 分类器模型,对检测到的目标进行分类。
# Set Dataloader
# 这段代码是 detect 函数中用于设置数据加载器(Dataloader)的部分。它根据输入源是网络摄像头还是文件(图像或视频)来决定使用哪种方式加载数据。
# 变量初始化。 这里初始化了两个变量 vid_path 和 vid_writer ,分别用于存储视频文件的路径和视频写入对象。在代码开始时,这两个变量被设置为 None 。
vid_path, vid_writer = None, None
# 网络摄像头检测。
# 这个条件判断基于之前定义的 webcam 变量。如果 webcam 为 True ,表示输入源是网络摄像头。
if webcam:
# 调用 check_imshow 函数来检查是否可以在屏幕上显示图像。这个函数可能返回 True 或 False ,取决于系统是否支持图像显示。
# def check_imshow():
# -> 检查当前环境是否支持图像显示,特别是使用 OpenCV 的 cv2.imshow() 函数和 PIL(Python Imaging Library)的 Image.show() 函数。如果环境支持图像显示,函数返回 True 。如果环境不支持图像显示,函数返回 False 。
# -> return True/return False
view_img = check_imshow()
# cudnn.benchmark
# cudnn.benchmark 是 PyTorch 中的一个设置,用于控制 NVIDIA 的 cuDNN 库是否在程序运行时自动为每个卷积层选择最优的算法。这个设置可以影响程序的性能,尤其是在深度学习模型中使用卷积层时。
# 定义和用法 :
# torch.backends.cudnn.benchmark :这是一个布尔值设置,可以设置为 True 或 False 。
# True :开启 cuDNN 的基准测试模式。在这个模式下,cuDNN 会在程序开始运行时为每个卷积层自动选择最优的算法。这可能会在程序启动时增加一些额外的时间开销,因为 cuDNN 需要对不同的算法进行基准测试,但一旦选择了最优算法,后续的卷积操作将会更快。
# False :关闭基准测试模式。cuDNN 将使用默认的卷积算法,这可能不是最优的选择,但适用于模型输入尺寸在运行过程中会改变的情况。
# 适用场景 :
# 固定输入尺寸 :如果你的模型输入尺寸(例如,图像尺寸和批处理大小)是固定的,设置 torch.backends.cudnn.benchmark = True 可以提高运行效率,因为 cuDNN 可以预先选择最优算法。
# 变化输入尺寸 :如果输入尺寸可能发生变化,开启 benchmark 可能导致性能下降,因为每次输入尺寸改变时,cuDNN 都可能重新搜索算法。
# 注意事项 :
# 性能影响 :开启 cudnn.benchmark 可能会在程序启动时增加一些额外的时间开销,但可以提高后续卷积操作的速度。
# 结果可重复性 :开启 cudnn.benchmark 可能会导致结果的轻微变化,因为 cuDNN 可能会选择不同的算法。如果需要确保结果的完全可重复性,可能需要关闭 cudnn.benchmark 并设置 torch.backends.cudnn.deterministic = True 。
# 总的来说, cudnn.benchmark 是一个有用的设置,可以帮助优化深度学习模型的性能,但需要根据具体的应用场景和需求来决定是否开启。
# 设置 CUDA 深度神经网络库(cuDNN)的基准测试为 True 。这可以加速推理过程,特别是当输入图像尺寸保持不变时。
cudnn.benchmark = True # set True to speed up constant image size inference
# 创建一个 LoadStreams 数据集对象,用于从网络摄像头或视频流中加载数据。 source 是视频流的地址或摄像头索引, img_size 是图像尺寸, stride 是模型的步长。
# class LoadStreams: -> 它用于从多个 IP 摄像头或 RTSP 流加载视频流。 -> def __init__(self, sources='streams.txt', img_size=640, stride=32):
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
# 文件检测。
# 部分代码执行的条件是输入源不是网络摄像头,即从文件(图像或视频)中加载数据。
else:
# 创建一个 LoadImages 数据集对象,用于从文件系统中加载图像或视频文件。 source 是图像或视频文件的路径, img_size 是图像尺寸, stride 是模型的步长。
# class LoadImages: -> LoadImages 的类,它用于在推理(inference)阶段加载图像或视频文件。 -> def __init__(self, path, img_size=640, stride=32):
dataset = LoadImages(source, img_size=imgsz, stride=stride)
# 这段代码根据输入源的不同,设置了不同的数据加载器。对于网络摄像头或视频流,使用 LoadStreams ;对于图像或视频文件,使用 LoadImages 。这样的设计使得 detect 函数能够灵活地处理不同类型的输入源,无论是实时的网络摄像头流还是静态的图像/视频文件。
# Get names and colors
# 这段代码是 detect 函数中用于获取模型类别名称和为每个类别生成随机颜色的部分。
# 获取类别名称。
# 这行代码尝试从模型中获取类别名称。 hasattr(model, 'module') 检查模型对象是否有一个名为 module 的属性,这通常出现在模型被包装在另一个类中的情况下。
# 如果 model 有 module 属性,那么类别名称将从 model.module.names 获取;如果没有,直接从 model.names 获取。
names = model.module.names if hasattr(model, 'module') else model.names
# 生成随机颜色。
# 这行代码为每个类别生成一个随机颜色。颜色由三个随机整数组成,分别代表 RGB 颜色空间中的红色、绿色和蓝色分量。每个分量的值在 0 到 255 之间。
# [[random.randint(0, 255) for _ in range(3)] for _ in names] 这个列表推导式首先为每个类别名称生成一个包含三个随机整数的列表,这三个整数分别对应 RGB 值。
# 然后,为 names 列表中的每个类别名称执行这个操作,最终生成一个颜色列表,其中每个颜色对应一个类别。
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
# 这段代码的作用是为检测到的每个类别准备一个唯一的颜色,以便在可视化检测结果时,不同的类别可以用不同的颜色表示。这对于在图像或视频上绘制边界框和标签时区分不同类别非常有用。
# Run inference
# 这段代码是 detect 函数中用于执行模型推理的部分。
# 模型预热。
if device.type != 'cpu':
# 这行代码检查是否使用的是非CPU设备(如GPU)。如果是,它会在模型上运行一次推理,以预热设备。
# 这是通过传递一个全零的张量给模型实现的,张量的大小为 (1, 3, imgsz, imgsz) ,其中 imgsz 是图像的尺寸, 3 代表颜色通道数(RGB)。这个张量被发送到指定的设备,并转换为模型参数的数据类型。
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
# 计时开始。这行代码记录了推理开始的时间,以便后续计算推理过程的总耗时。
t0 = time.time()
# 遍历数据集。
# 这行代码开始遍历数据集, dataset 是之前创建的 LoadStreams 或 LoadImages 数据集对象。对于每个数据项,它返回 路径 path 、 图像 img 、 原始图像 im0s 和 视频捕获对象 vid_cap (如果是视频流)。
for path, img, im0s, vid_cap in dataset:
# 图像预处理。
# 将从数据集中获取的图像 img 从 NumPy 数组转换为 PyTorch 张量,并发送到指定的设备。
img = torch.from_numpy(img).to(device)
# 如果 half 为 True (即设备支持半精度计算),则将图像张量转换为半精度浮点数(FP16)。否则,转换为单精度浮点数(FP32)。这一步还将图像的数据类型从无符号整型(uint8)转换为浮点数。
img = img.half() if half else img.float() # uint8 to fp16/32
# 将图像的像素值从 [0, 255] 范围归一化到 [0.0, 1.0] 范围。
img /= 255.0 # 0 - 255 to 0.0 - 1.0
# torch.Tensor.ndimension()
# 在PyTorch中, ndimension() 方法用于返回一个张量(Tensor)的维度数目。这个方法是 torch.Tensor 类的一个实例方法,可以用于确定张量有多少个维度。
# 返回值 :
# 返回一个整数,表示张量的维度数。
# 这个方法在处理多维数据时非常有用,尤其是在需要根据张量的维度进行条件判断或循环操作时。例如,你可以使用 ndimension() 来检查一个张量是否是标量(0维)、向量(1维)或矩阵(2维),并据此执行不同的操作。
# 如果图像张量只有三个维度(即没有批次维度),则添加一个批次维度。这是因为模型通常期望输入数据有一个批次维度。
if img.ndimension() == 3:
img = img.unsqueeze(0)
# 这段代码负责准备模型进行推理,包括模型预热、遍历数据集、图像预处理和添加批次维度。这些步骤是模型推理前的标准流程,确保输入数据符合模型的期望格式。
# 这段代码是 detect 函数中用于执行 模型推理 、 应用非极大值抑制(NMS) 和 可选的分类器 的部分。
# Inference
# 同步计时开始。 time_synchronized() 是一个自定义函数,用于获取一个精确的时间戳,用于同步不同设备上的时间(例如,在多GPU环境中)。这里记录了推理开始的时间。
# def time_synchronized(): -> 它用于获取一个与CUDA时间同步的准确时间戳。 -> return time.time()
t1 = time_synchronized()
# 模型推理。
# 这行代码将预处理后的图像 img 传递给模型进行推理。 opt.augment 是一个布尔值,指示是否对图像进行数据增强。模型返回的结果是一个包含预测结果的张量, [0] 表示我们取第一个元素,即原始尺寸的预测结果。
pred = model(img, augment=opt.augment)[0]
# Apply NMS
# 应用非极大值抑制(NMS)。
# non_max_suppression 是一个函数,用于抑制重叠的检测框,只保留最佳的检测结果。
# opt.conf_thres 是置信度阈值, opt.iou_thres 是交并比阈值,用于确定何时抑制重叠的检测框。 classes=opt.classes 指定了要保留的类别。 agnostic=opt.agnostic_nms 表示是否使用与类别无关的NMS。
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
# 同步计时结束。再次调用 time_synchronized() 记录推理和NMS之后的时间。
# def time_synchronized(): -> 它用于获取一个与CUDA时间同步的准确时间戳。 -> return time.time()
t2 = time_synchronized()
# Apply Classifier
# 应用分类器(如果启用)。
if classify:
# 如果 classify 为 True ,则调用 apply_classifier 函数,将分类器应用于模型的预测结果。 pred 是经过NMS后的预测结果, modelc 是分类器模型, img 是输入图像的张量, im0s 是原始输入图像。
# def apply_classifier(x, model, img, im0): -> 用于对 YOLO(You Only Look Once)目标检测模型的输出应用第二阶段分类器。回处理后的检测结果。 -> return x
pred = apply_classifier(pred, modelc, img, im0s)
# 这段代码负责执行模型的推理过程,并通过NMS过滤掉重叠的检测框,以保留最佳的检测结果。如果启用了分类器,还会对检测到的目标进行进一步的分类。这些步骤是目标检测流程中的关键部分,确保了检测结果的准确性和可靠性。
# Process detections
# 这段代码是 detect 函数中用于处理每个图像的检测结果的部分。
# 遍历每个图像的检测结果。 这行代码遍历 pred 列表中的每个元素, pred 是包含所有图像检测结果的列表。 i 是索引, det 是对应索引的检测结果。
for i, det in enumerate(pred): # detections per image
# 处理网络摄像头和非网络摄像头的情况。
if webcam: # batch_size >= 1
# 如果是网络摄像头( webcam 为 True ),则从 path 、 im0s 等列表中获取第 i 个元素,因为网络摄像头会产生批次(batch)数据。
# 路径 ( path ) :这里, path 是一个包含图像路径或视频流信息的列表。通过索引 i 访问这个列表,获取当前处理的图像或视频流的具体路径或信息。
# 字符串标签 ( s ) :这里创建了一个字符串 s ,它以当前索引 i 开始,通常用于后续打印或记录信息时标记当前处理的是批次中的哪一张图像。
# 图像副本 ( im0 ) : im0s 是一个包含原始图像的列表。通过索引 i 获取当前图像的副本,并使用 .copy() 方法确保获得原始图像数据的一个独立副本,这样可以在不修改原始数据的情况下进行操作。
# 帧编号 ( frame ) :dataset.count 是一个计数器,用于跟踪当前处理的是视频流中的第几帧。这对于视频流处理特别有用,因为它们通常以帧为单位连续输入。
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
else:
# 如果不是网络摄像头,那么处理的是单张图像或视频帧,所以不需要索引 i 。
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
# 路径和文件名处理。
# 将路径 p 转换为 Path 对象,以便使用 pathlib 库进行路径操作。
p = Path(p) # to Path
# 构建保存检测结果图像的路径。
save_path = str(save_dir / p.name) # img.jpg
# p.stem :返回文件名不包括后缀的部分。
# 构建保存检测结果文本文件的路径。如果数据集模式是图像,则直接使用 p.stem (即文件名不带扩展名);如果是视频,则添加帧编号。
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
# 构建打印字符串和归一化增益。
# 构建一个字符串 s ,包含图像的尺寸。
# s :这是一个字符串变量,用于累加图像信息。
# '%gx%g ' :这是一个格式化字符串,包含两个 %g 占位符,用于插入两个数值。
# 格式化过程 :第一个 %g 被替换为 img.shape[2] 的值,即图像的高度。 第二个 %g 被替换为 img.shape[3] 的值,即图像的宽度。
# 结果 :假设 img 是一个形状为 (1, 3, 640, 480) 的张量(例如,一个批处理大小为1,3个颜色通道,宽度为640像素,高度为480像素的图像),这行代码将把字符串 s 更新为: s = '640x480 ' 。
s += '%gx%g ' % img.shape[2:] # print string
# 创建一个张量 gn ,包含原始图像 im0 的宽度和高度,用于后续的坐标归一化。
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
# 处理检测结果。
if len(det):
# Rescale boxes from img_size to im0 size
# 如果检测结果 det 不为空,那么需要将检测框的坐标从模型输入图像的大小( img_size )重新缩放到原始图像的大小( im0 )。 scale_coords 函数用于执行这个缩放操作,并且结果被四舍五入到最近的整数。
# def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): -> 它用于将坐标从一张图片的形状 ( img1_shape ) 缩放到另一张图片的形状 ( img0_shape )。返回调整后的坐标。 -> return coords
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# 这段代码负责处理每个图像的检测结果,包括路径和文件名的构建、打印字符串的构建、归一化增益的计算以及检测框坐标的重新缩放。这些步骤是将模型输出的检测结果转换为可用于显示和保存的格式的关键过程。
# 这段代码是 detect 函数中用于打印检测结果和将结果写入文件或图像的部分。
# Print results
# 打印检测结果。
# 这段代码遍历 det 中所有唯一的类别 c 。 det[:, -1] 获取检测结果中的最后一列,即类别索引。 unique() 函数获取这些类别索引的唯一值。
for c in det[:, -1].unique():
# 对于每个唯一的类别 c ,计算该类别的检测数量 n 。
n = (det[:, -1] == c).sum() # detections per class
# 然后将检测数量和类别名称添加到字符串 s 中,如果检测数量大于1,则在类别名称后加上 's' 。
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
# 写入结果到文件。
# reversed(seq)
# reversed() 是 Python 内置的一个函数,用于返回一个反向迭代器。它可以对任何可迭代对象(如列表、元组、字符串等)进行反向迭代。
# 参数 :
# seq :这是要反转的可迭代对象。可以是列表、元组、字符串等。
# 返回值 :
# 返回一个反向迭代器,可以通过 for 循环或其他迭代方式访问。
# 注意事项 :
# reversed() 不会修改原始序列,而是返回一个新的迭代器。
# 如果传入的对象不支持反向迭代(例如,整数),将会引发 TypeError 。
# 总结 :
# reversed() 是一个非常实用的函数,特别是在需要以相反的顺序处理可迭代对象时。它提供了一种简单而有效的方法来反转迭代顺序,而不需要手动创建新的列表或其他数据结构。
# 这段代码遍历 det 中的每个检测结果, xyxy 表示边界框的坐标, conf 表示置信度, cls 表示类别索引。
for *xyxy, conf, cls in reversed(det):
# 如果 save_txt 为 True ,则将检测结果写入文本文件。
if save_txt: # Write to file
# xyxy2xywh 函数将边界框的坐标从 xyxy 格式转换为 xywh 格式(即从左上角和右下角坐标转换为中心点坐标和宽高)。
# gn 是归一化增益,用于将坐标从模型输入图像的大小缩放到原始图像的大小。
# def xyxy2xywh(x): -> 用于将边界框的坐标从 xyxy 格式(即左上角和右下角的坐标)转换为 xywh 格式(即中心点坐标加上宽度和高度)。返回结果。返回转换后的 xywh 格式的边界框坐标。 -> return y
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# line 是要写入文件的数据,包括类别索引、归一化的 xywh 坐标和置信度(如果 opt.save_conf 为 True )。
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
# 使用 with open 语句打开文件,并以追加模式写入数据。
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 在图像上绘制边界框。
# 如果 save_img 或 view_img 为 True ,则在图像 im0 上绘制边界框。
if save_img or view_img: # Add bbox to image
# label 是要显示在边界框旁边的文本,包括类别名称和置信度。
label = f'{names[int(cls)]} {conf:.2f}'
# plot_one_box 函数在图像上绘制一个边界框, xyxy 是边界框的坐标, label 是标签文本, color 是边界框的颜色, line_thickness 是边界框线条的厚度。
# def plot_one_box(x, img, color=None, label=None, line_thickness=3): -> 用于在图像上绘制单个边界框,并可选地添加标签。这个函数通常用于目标检测任务中,用于在检测到的对象周围绘制矩形框,并显示类别名称和置信度。
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
# 这段代码负责将检测结果以文本形式打印出来,并根据配置将结果写入文件或在图像上绘制边界框。这是目标检测流程中展示和记录检测结果的关键步骤。
# 这段代码是 detect 函数中用于打印推理时间、显示图像和保存检测结果的部分。
# 打印推理时间(可选)。
# Print time (inference + NMS)
# 这一行被注释掉了,但如果取消注释,它会打印出处理当前图像(包括推理和NMS)所花费的时间。 t2 - t1 计算了从推理开始到结束的时间差。
#print(f'{s}Done. ({t2 - t1:.3f}s)')
# Stream results
# 显示图像。
if view_img:
# 如果 view_img 为 True ,则使用 cv2.imshow 显示图像 im0 。 str(p) 作为窗口标题, im0 是要显示的图像。
cv2.imshow(str(p), im0)
# cv2.waitKey(1) 使图像显示窗口等待1毫秒,以便捕获按键事件,例如可以按任意键关闭窗口。
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections) 保存结果(带有检测的图像)。
# 保存检测结果图像。如果 save_img 为 True ,则执行以下操作来保存图像或视频。
if save_img:
# 保存单张图像。
if dataset.mode == 'image':
# 如果数据集模式是 'image' ,则使用 cv2.imwrite 将带有检测结果的图像 im0 保存到 save_path 指定的路径。
cv2.imwrite(save_path, im0)
# 保存视频或流。
# 如果数据集模式不是 'image' ,则处理视频或流。
else: # 'video' or 'stream'
# 检查新视频路径。这行代码检查当前视频路径 vid_path 是否与要保存的路径 save_path 不同。如果不同,表示正在处理一个新的视频文件。
if vid_path != save_path: # new video
# 更新视频路径。更新 vid_path 变量为新的保存路径 save_path 。
vid_path = save_path
# 释放旧的视频写入对象。
if isinstance(vid_writer, cv2.VideoWriter):
# 如果 vid_writer 是 cv2.VideoWriter 类型,调用 release() 方法释放之前的 VideoWriter 对象,以便创建新的 VideoWriter 对象。
vid_writer.release() # release previous video writer
# 获取视频属性。
if vid_cap: # video
# 如果 vid_cap 存在(即处理的是视频文件),则从视频捕获对象中获取 帧率(FPS) 、 宽度 和 高度 。
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
# 如果不是处理视频文件(即处理的是视频流),则设置默认的帧率为30,宽度和高度从图像 im0 的尺寸获取,并确保 save_path 包含文件扩展名 .mp4 。
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
# cv2.VideoWriter(filename, fourcc, fps, frameSize[, isColor=1])
# cv2.VideoWriter 是 OpenCV 库中的一个类,用于将视频帧写入视频文件。
# 参数 :
# filename :输出视频文件的名称或路径。
# fourcc :一个四字符代码(Four-Character Code),指定视频编解码器。常见的编解码器包括 'XVID' 、 'MJPG' 、 'X264' 和 'mp4v' 等。
# fps :视频的帧率(frames per second),即每秒的帧数。
# frameSize :视频帧的尺寸,通常以 (width, height) 的形式给出。
# isColor :(可选)一个布尔值,指示输入帧是否为彩色。默认值为 1 ,表示彩色;如果设置为 0 ,则表示灰度图像。
# 返回值 :
# 返回一个 cv2.VideoWriter 对象,用于后续的帧写入操作。
# 注意事项 :
# 在写入完所有帧后,应该调用 release() 方法来释放 VideoWriter 对象,确保所有帧都被写入文件。
# cv2.VideoWriter 支持多种视频格式和编解码器,但具体支持哪些取决于你的系统和 OpenCV 配置。
# 如果视频帧的尺寸在写入过程中发生变化,可能需要重新创建 VideoWriter 对象,因为 frameSize 必须与实际写入的帧尺寸匹配。
# cv2.VideoWriter 是处理视频文件输出的强大工具,可以用于保存视频流、录制摄像头视频、保存检测结果等场景。
# cv2.VideoWriter 类是 OpenCV 中用于创建和写入视频文件的工具。以下是 cv2.VideoWriter 类中包含的一些常用方法 :
# 1. 构造函数
# VideoWriter() :默认构造函数,创建一个未初始化的 VideoWriter 对象。
# VideoWriter(const String& filename, int fourcc, double fps, Size frameSize, bool isColor=true) :初始化 VideoWriter 对象,设置输出视频文件的名称、编解码器、帧率、帧尺寸和颜色模式。
# 2. open()
# open(const String& filename, int fourcc, double fps, Size frameSize, bool isColor=true) :打开或重新初始化视频写入器。
# 3. write()
# write(const Mat& image) :写入下一帧图像到视频文件中。也可以使用重载的 << 操作符。
# 4. release()
# release() :释放 VideoWriter 对象,关闭视频文件。
# 5. set()
# set(int propId, double value) :设置视频写入器的属性。
# 6. get()
# get(int propId) const :获取视频写入器的属性值。
# 7. fourcc()
# static int fourcc(char c1, char c2, char c3, char c4) :创建一个四字符代码,用于指定视频编解码器。
# 8. isOpened()
# isOpened() const :检查 VideoWriter 对象是否已成功初始化。
# 9. getBackendName()
# getBackendName() const :返回使用的后端 API 名称。
# 这些方法提供了创建、配置、写入帧和关闭视频文件所需的全部功能。使用 cv2.VideoWriter 类时,通常首先创建一个对象,然后使用 write() 方法写入帧,最后调用 release() 方法来关闭文件。
# 创建新的视频写入对象。创建一个新的 cv2.VideoWriter 对象,用于写入视频帧。 save_path 是视频文件的保存路径, cv2.VideoWriter_fourcc(*'mp4v') 是视频编解码器, fps 是帧率, (w, h) 是视频的宽度和高度。
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
# 写入帧数据。使用 vid_writer.write(im0) 将图像 im0 (可能包含检测结果)写入视频文件。
vid_writer.write(im0)
# 这段代码负责在处理新视频时正确地初始化视频写入对象,并确保旧的视频写入对象被释放。然后,它将图像帧写入视频文件,这对于保存检测结果视频非常有用。
# 这段代码是 detect 函数中用于在检测过程结束后打印结果和统计信息的部分。
# 条件检查。这行代码检查是否需要保存文本文件( save_txt )或图像文件( save_img )。如果任一条件为真,则执行以下代码块。
if save_txt or save_img:
# 统计和打印保存的标签。
# 如果 save_txt 为 True ,则计算 save_dir 下 labels 子目录中 .txt 文件的数量,并构建一个字符串 s ,包含保存的标签数量和保存路径。
# 如果 save_txt 为 False ,则 s 为空字符串。
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
# 打印结果保存位置(可选)。这一行被注释掉了,但如果取消注释,它会打印出结果保存的位置,包括保存目录和保存的标签数量信息。
#print(f"Results saved to {save_dir}{s}")
# 打印总耗时。这行代码计算从 t0 (检测开始时间)到现在的总耗时,并打印出来。 time.time() 返回当前时间的时间戳, t0 是检测开始时记录的时间戳。
print(f'Done. ({time.time() - t0:.3f}s)')
# 这段代码在检测过程结束后提供反馈,包括结果保存的位置和总耗时。这对于用户了解检测过程的效率和结果非常有用。通过打印出保存的标签数量和路径,用户可以方便地找到和查看检测结果。同时,总耗时的打印有助于评估检测过程的性能。
# 这个函数实现了一个完整的图像或视频检测流程,包括模型加载、图像处理、推理、结果处理和保存。通过这个函数,可以实现对图像或视频的实时检测,并保存检测结果。
3.if __name__ == '__main__':
# if __name__ == '__main__' : 是一个常用的模式,用于判断当前脚本是否作为主程序运行。如果是,那么以下代码块将被执行;如果不是(即该模块被其他 Python 脚本导入),则不执行。
if __name__ == '__main__':
# argparse.ArgumentParser() 是 Python 标准库 argparse 模块中的一个函数,用于创建命令行参数解析器。这个解析器可以定义需要哪些命令行参数,每个参数的类型、默认值、帮助信息等。
parser = argparse.ArgumentParser()
# --weights : 模型权重文件的路径,可以是单个文件或多个文件。
parser.add_argument('--weights', nargs='+', type=str, default='yolov7.pt', help='model.pt path(s)')
# --source : 指定检测的输入源,可以是文件、文件夹或摄像头(0表示摄像头)。
parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
# --img-size : 用于推理的图像尺寸(像素)。
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
# --conf-thres : 目标置信度阈值。
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
# --iou-thres : 用于非极大值抑制(NMS)的交并比(IOU)阈值。
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
# --device : 指定使用的设备,例如 GPU(如 '0' 或 '0,1,2,3')或 CPU。
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# --view-img : 是否显示结果。
parser.add_argument('--view-img', action='store_true', help='display results')
# --save-txt : 是否将结果保存到文本文件。
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
# --save-conf : 是否在保存的文本标签中保存置信度。
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
# --nosave : 是否不保存图像或视频。
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
# --classes : 通过类别过滤检测结果。
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
# --agnostic-nms : 是否使用类别不可知的 NMS。
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
# --augment : 是否进行增强推理。
parser.add_argument('--augment', action='store_true', help='augmented inference')
# --update : 是否更新所有模型(用于修复 SourceChangeWarning)。
parser.add_argument('--update', action='store_true', help='update all models')
# --project : 保存结果的项目路径。
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
# --name : 保存结果的名称。
parser.add_argument('--name', default='exp', help='save results to project/name')
# --exist-ok : 是否允许项目/名称已存在,不进行增量。
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
# --trace : 是否追踪模型。
parser.add_argument('--trace', action='store_true', help='trace model')
# 解析命令行参数。
opt = parser.parse_args()
print(opt)
#check_requirements(exclude=('pycocotools', 'thop'))
# with torch.no_grad(): 是 PyTorch 中的一个上下文管理器,用于指示 PyTorch 在代码块执行期间不要跟踪梯度。这通常用在推理(inference)阶段,因为推理不需要进行反向传播,因此不需要计算梯度,这样可以减少内存消耗并加快计算速度。
# detect() 函数涉及到模型的前向传播(例如,进行预测),将 with torch.no_grad(): 包裹在调用 detect() 函数的代码块周围是合适的。这样可以确保在推理过程中不计算梯度。
with torch.no_grad():
# 这段代码中的 if opt.update: 块有一个逻辑错误。它试图在列表中迭代 opt.weights ,但 opt.weights 是一个字符串,不是列表。正确的做法应该是将 opt.weights 定义为列表,然后在循环中迭代。例如:
if opt.update: # update all models (to fix SourceChangeWarning)
for opt.weights in ['yolov7.pt']:
detect()
# def strip_optimizer(f='best.pt', s=''): -> 目的是从训练好的模型文件中移除优化器(optimizer)和其他非必要的信息,以便减小模型文件的大小,使其更适合部署。这个函数还允许你将处理后的模型保存为一个新的文件。
strip_optimizer(opt.weights)
else:
# 如果不需要更新模型,直接调用 detect 函数。
detect()
标签:YOLOv7,opt,detect,img,0.1,模型,--,图像,save
From: https://blog.csdn.net/m0_58169876/article/details/143859417