inferer.py
yolov6\core\inferer.py
目录
inferer.py
1.所需的库和模块
2.class Inferer:
3.class CalcFPS:
1.所需的库和模块
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# 用于模型的推理。
import os
import cv2
import time
import math
import torch
import numpy as np
import os.path as ospfrom tqdm import tqdm
from pathlib import Path
from PIL import ImageFont
from collections import dequefrom yolov6.utils.events import LOGGER, load_yaml
from yolov6.layers.common import DetectBackend
from yolov6.data.data_augment import letterbox
from yolov6.data.datasets import LoadData
from yolov6.utils.nms import non_max_suppression
from yolov6.utils.torch_utils import get_model_info
2.class Inferer:
class Inferer:def __init__(self, source, webcam, webcam_addr, weights, device, yaml, img_size, half):# 将构造函数的参数直接添加到实例属性中。self.__dict__.update(locals())# Init model 初始化模型。# 确保设备属性被正确设置。self.device = device# 确保图像尺寸属性被正确设置。self.img_size = img_size# 使用 PyTorch 的 torch.device 来设置模型运行的设备。如果 cuda 为 True （即用户指定了 GPU 设备并且 CUDA 可用），则设备设置为 cuda:device ，其中 device 是传入的参数。否则，设备设置为 'cpu' 。cuda = self.device != 'cpu' and torch.cuda.is_available()# torch.device(device_str)# torch.device 是 PyTorch 中的一个函数，用于创建一个表示设备的对象，该对象可以是 CPU 或者 CUDA（GPU）。这个函数的主要作用是抽象出设备的选择，使得代码可以在不同的设备上运行而不需要做太多修改。# 其中 device_str 是一个字符串，可以是以下几种形式 ：# 1. 'cpu' ：表示使用 CPU。# 2. 'cuda' ：表示使用第一个 CUDA 设备（GPU）。# 3. 'cuda:N' ：表示使用编号为 N 的 CUDA 设备，其中 N 是一个非负整数。# 4. 'cuda:0' ：表示使用编号为 0 的 CUDA 设备。# 5. 'cuda:1' ：表示使用编号为 1 的 CUDA 设备，以此类推。# 如果指定了 CUDA 设备，但是机器上没有可用的 CUDA 设备， torch.device 会抛出一个错误。# 除了指定设备， torch.device 还可以接受一个字典或者一个 torch.device 对象作为参数，这时候它会返回一个与输入相同的 torch.device 对象。self.device = torch.device(f'cuda:{device}' if cuda else 'cpu')# 使用 DetectBackend 类来加载模型。 weights 参数是模型权重文件的路径， device 是模型将要运行的设备。 DetectBackend 类可能是 YOLOv6 框架中用于模型推理的后端。# class DetectBackend(nn.Module): -> def __init__(self, weights='yolov6s.pt', device=None, dnn=True):self.model = DetectBackend(weights, device=self.device)# 从模型中获取步长（ stride ），这个属性通常用于图像尺寸的调整和锚点（anchor）的计算。self.stride = self.model.stride# def load_yaml(file_path): -> 从 yaml 文件加载数据。函数返回从 YAML 文件中加载的数据，通常是一个字典或列表，具体取决于 YAML 文件的内容。 -> return data_dicts# 使用 load_yaml 函数从 YAML 文件中加载类别名称。这个 YAML 文件包含了数据集的类别信息， ['names'] 键对应的值是一个包含所有类别名称的列表。self.class_names = load_yaml(yaml)['names']# 调用 self.check_img_size 方法来检查和调整 img_size ，以确保它与模型的步长兼容。这个方法可能是自定义的，用于确保输入图像的尺寸是模型所支持的。# def check_img_size(self, img_size, s=32, floor=0):# -> 作用是确保输入的图像尺寸 img_size 是模型步长 s 的倍数。如果 img_size 是列表，直接返回调整后的尺寸列表。如果 img_size 是整数，返回一个包含两个相同值的列表，这两个值都是调整后的尺寸。# -> return new_size if isinstance(img_size,list) else [new_size]*2self.img_size = self.check_img_size(self.img_size, s=self.stride) # check image size 检查图像大小。# 将 half 属性设置为传入的值，这个属性可能用于控制是否使用半精度（FP16）进行推理，以加速推理过程并减少内存使用。self.half = half# 切换模型至部署状态。# Switch model to deploy statusself.model_switch(self.model.model, self.img_size)# Half precision# self.half ：一个布尔值，指示是否使用半精度（FP16）。# self.device.type != 'cpu' ：确保设备不是CPU，因为CPU不支持半精度运算。if self.half & (self.device.type != 'cpu'):# 如果条件为真，则调用 self.model.model.half() 将模型转换为半精度。self.model.model.half()else:# 否则，调用 self.model.model.float() 将模型转换为全精度（FP32），并将 self.half 设置为 False 。self.model.model.float()self.half = False# 这个条件确保只在非CPU设备上执行预热操作。if self.device.type != 'cpu':# 这一行执行一次模型推理，使用一个全零的张量作为输入 ：# torch.zeros(1, 3, *self.img_size) ：创建一个形状为 (1, 3, img_size, img_size) 的全零张量，其中 1 是批量大小， 3 是颜色通道数（对于RGB图像）， img_size 是图像的尺寸。# .to(self.device) ：将张量移动到指定的设备（如GPU）。# .type_as(next(self.model.model.parameters())) ：确保输入张量的数据类型与模型参数的数据类型一致。self.model(torch.zeros(1, 3, *self.img_size).to(self.device).type_as(next(self.model.model.parameters()))) # warmup# 加载数据。# Load data# 一个布尔值，指示是否使用网络摄像头作为输入源。self.webcam = webcam# 如果使用网络摄像头，这是摄像头的地址。self.webcam_addr = webcam_addr# 调用 LoadData 类（可能是自定义的）来加载数据。这个类根据提供的 source 、 webcam 和 webcam_addr 参数确定数据来源，并返回一个包含数据的实例。# 如果 webcam 为 True ，则 LoadData 可能会连接到网络摄像头并开始捕获视频流。# 如果 source 是文件路径或文件夹， LoadData 可能会读取文件列表或视频帧。self.files = LoadData(source, webcam, webcam_addr)# 输入源的路径，可以是单个文件、文件夹或摄像头。self.source = source# 将模型中的 RepVGGBlock 层切换到部署（deploy）状态。 RepVGGBlock 是一个可重参数化的网络层，它在训练和部署时有不同的结构。在训练时， RepVGGBlock 会使用一种结构以便于学习，而在部署时，它会切换到另一种更高效的结构以减少计算量和模型大小。# 1.self ：类的实例本身。# 2.model ：需要被切换到部署状态的模型。# 3.img_size ：图像的尺寸，尽管在这个函数中并没有直接使用。def model_switch(self, model, img_size):# 模型切换到部署状态。''' Model switch to deploy status '''# 从 yolov6.layers.common 模块导入 RepVGGBlock 类。from yolov6.layers.common import RepVGGBlock# 使用 model.modules() 遍历模型中的所有层（模块）。for layer in model.modules():# 对于每个层，检查它是否是 RepVGGBlock 的实例。if isinstance(layer, RepVGGBlock):# 如果层是 RepVGGBlock 的实例，调用该层的 switch_to_deploy 方法将其切换到部署状态。# def switch_to_deploy(self):layer.switch_to_deploy()# 切换模型以部署模式。LOGGER.info("Switch model to deploy modality.")# conf_thres ：置信度阈值，用于过滤检测结果。# iou_thres ：交并比（Intersection over Union）阈值，用于非极大值抑制（NMS）。# classes ：要检测的类别列表。# agnostic_nms ：是否使用与类别无关的NMS。# max_det ：每张图像最大检测数量。# save_dir ：保存结果的目录。# save_txt ：是否保存检测结果的文本文件。# save_img ：是否保存带有检测框的图像。# hide_labels ：是否隐藏标签。# hide_conf ：是否隐藏置信度分数。# view_img ：是否在推理时显示图像。def infer(self, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf, view_img=True):# 模型推理和结果可视化。''' Model Inference and results visualization '''# vid_path ：用于存储视频路径，初始设置为 None 。# vid_writer ：用于存储视频写入对象，初始设置为 None 。这可能用于将检测结果写入视频文件。# windows ：用于存储窗口信息的列表，初始为空列表。这可能用于处理多窗口或多图像的情况。vid_path, vid_writer, windows = None, None, []# fps_calculator ：创建一个 CalcFPS 类的实例，用于计算帧率。# class CalcFPS: -> def __init__(self, nsamples: int = 50):fps_calculator = CalcFPS()# 使用 tqdm 进度条遍历 self.files ，这是一个由 LoadData 类创建的数据迭代器。# img_src ：当前帧图像。# img_path ：图像的文件路径。# vid_cap ：视频捕获对象，用于从视频文件中读取帧。for img_src, img_path, vid_cap in tqdm(self.files):# 调用 process_image 方法对图像进行预处理，包括缩放、归一化等操作。# def process_image(img_src, img_size, stride, half): -> 在图像推理之前处理图像。返回值 ： 1.image ：预处理后的图像张量。 2.img_src ：原始图像数据或路径。 -> return image, img_srcimg, img_src = self.process_image(img_src, self.img_size, self.stride, self.half)# 将预处理后的图像移动到指定的设备（如GPU）。img = img.to(self.device)# 检查图像的维度，如果只有3维（通道、高度、宽度），则添加一个批次维度。if len(img.shape) == 3:# 通过 None 索引添加批次维度。img = img[None]# expand for batch dim 扩展批量维度。# 记录推理开始的时间。t1 = time.time()# 将预处理后的图像输入模型，获取预测结果。pred_results = self.model(img)# 调用 non_max_suppression 函数对预测结果应用NMS，过滤掉低置信度的检测框，并去除重叠的检测框。# 1.pred_results ：模型的预测结果。# 2.conf_thres ：置信度阈值。# 3.iou_thres ：交并比阈值。# 4.classes ：要检测的类别列表。# 5.agnostic_nms ：是否使用与类别无关的NMS。# 6.max_det ：每张图像最大检测数量。det = non_max_suppression(pred_results, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)[0]# 记录推理结束的时间。t2 = time.time()# self.webcam ：一个布尔值，指示是否使用网络摄像头作为输入源。if self.webcam:# 如果 self.webcam 为 True ，则保存路径 save_path 和文本路径 txt_path 都设置为 save_dir 目录下，以 self.webcam_addr 为名称的子目录。save_path = osp.join(save_dir, self.webcam_addr)txt_path = osp.join(save_dir, self.webcam_addr)else:# 在嵌套目录中创建镜像图像目录结构的输出文件。# Create output files in nested dirs that mirrors the structure of the images' dirs# dirname = os.path.dirname(path)# osp.dirname 是 Python os.path 模块中的一个函数，它用于返回路径中的目录名称。# 如果给定的路径是一个文件路径， osp.dirname 将返回该文件所在的目录路径；如果给定的路径本身就是一个目录，它将返回该目录的父目录路径。# path ：输入的路径字符串。# 返回值 : dirname 是一个字符串，表示输入路径的目录部分。# 如果 self.webcam 为 False ，则根据输入图像的路径和输入源的路径来创建一个相对路径 rel_path 。rel_path = osp.relpath(osp.dirname(img_path), osp.dirname(self.source))# 保存图像结果的路径，设置为 save_dir 下，与输入图像相对位置相同的路径。save_path = osp.join(save_dir, rel_path, osp.basename(img_path)) # im.jpg# 保存检测结果文本文件的路径，设置为 save_dir 下，与输入图像相对位置相同的路径，但在 labels 子目录中，并去掉文件扩展名。txt_path = osp.join(save_dir, rel_path, 'labels', osp.splitext(osp.basename(img_path))[0])# 创建保存图像结果所需的目录结构，如果目录已存在，则不会抛出异常。os.makedirs(osp.join(save_dir, rel_path), exist_ok=True)# 这行代码创建了一个包含图像原始尺寸的 PyTorch 张量，然后通过索引 [1, 0, 1, 0] 来选择宽度和高度（重复两次）以形成 whwh 的格式。这种格式通常用于归一化图像尺寸，其中 w 是宽度， h 是高度。# 例如，如果 img_src 的形状是 (高度, 宽度, 通道数) ，那么 gn 将是一个包含 （宽度，高度，宽度，高度） 的张量。gn = torch.tensor(img_src.shape)[[1, 0, 1, 0]] # normalization gain whwh 归一化增益 whwh。# 这行代码创建了 img_src 图像的副本。这通常在需要保留原始图像数据进行后续操作或比较时使用。img_ori = img_src.copy()# 检查图像和字体。# check image and font# assert 是 Python 中的断言语句，用于验证某个条件是否为真。如果条件为假，则抛出 AssertionError 。# 这里检查 img_ori 的数据是否是连续的。在 PyTorch 中， Tensor 的内存布局可以是连续的或非连续的。连续的 Tensor 在内存中是一块连续的空间，这通常对于某些操作（如某些类型的复制或数学运算）是必需的。# 图像需要连续，请使用np.ascontiguousarray(im)来输入图像。assert img_ori.data.contiguous, 'Image needs to be contiguous. Please apply to input images with np.ascontiguousarray(im).'# 这个方法调用类实例 self 的一个方法，名为 font_check 。# def font_check(font='./yolov6/utils/Arial.ttf', size=10):# -> 用于检查指定路径的字体文件是否存在，并返回一个 PIL（Python Imaging Library）TrueType 字体对象。如果字体文件不存在，它将尝试下载字体并加载。# -> return ImageFont.truetype(str(font), size)self.font_check()# 检查检测结果 det 是否非空。if len(det):# 调整检测框的坐标以匹配原始图像的尺寸，并四舍五入到最近的整数。 self.rescale 方法用于将检测框从模型输出的坐标系转换回原始图像的坐标系。# def rescale(ori_shape, boxes, target_shape): -> 它用于将检测框的坐标从模型输入图像的尺寸重新缩放到原始图像的尺寸。重新缩放后的检测框坐标。 -> return boxesdet[:, :4] = self.rescale(img.shape[2:], det[:, :4], img_src.shape).round()# 遍历检测结果， xyxy 是检测框的坐标， conf 是置信度， cls 是类别索引。for *xyxy, conf, cls in reversed(det):# 检查是否需要保存检测结果到文本文件。if save_txt: # Write to file 写入文件。# 将检测框坐标从 xyxy 格式转换为 xywh 格式（中心点坐标和宽高），并归一化。# def box_convert(x):# -> 它用于将检测框的坐标从 [x1, y1, x2, y2] 格式（其中 (x1, y1) 是左上角坐标， (x2, y2) 是右下角坐标）转换为 [x, y, w, h] 格式（其中 (x, y) 是中心点坐标， w 是宽度， h 是高度）。y ：转换后的检测框坐标，格式为 [x, y, w, h] 。# -> return yxywh = (self.box_convert(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 归一化 xywh。# 创建一个包含 类别索引 、 归一化的 xywh 坐标 和 置信度 的元组。line = (cls, *xywh, conf)# 打开文本文件以追加模式。with open(txt_path + '.txt', 'a') as f:# 将检测结果写入文件。f.write(('%g ' * len(line)).rstrip() % line + '\n')# 检查是否需要在图像上绘制检测结果。if save_img:# 获取类别索引的整数值。class_num = int(cls) # integer class 整数类。# 根据是否隐藏标签和置信度，确定要绘制的标签文本。label = None if hide_labels else (self.class_names[class_num] if hide_conf else f'{self.class_names[class_num]} {conf:.2f}')# 在原始图像 img_ori 上绘制检测框和标签。# def plot_box_and_label(image, lw, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255), font=cv2.FONT_HERSHEY_COMPLEX):# -> 它用于在图像上绘制一个边界框（bounding box）以及与之关联的标签（label）。该方法没有返回值，它直接在输入的 image 上进行绘制。# def generate_colors(i, bgr=False):# -> 它用于生成一组预定义颜色的调色板，并根据给定的索引 i 返回相应的颜色。根据 bgr 参数返回 RGB 或 BGR 格式的颜色元组。# -> return (color[2], color[1], color[0]) if bgr else colorself.plot_box_and_label(img_ori, max(round(sum(img_ori.shape) / 2 * 0.003), 2), xyxy, label, color=self.generate_colors(class_num, True))# 将 PIL 图像 img_ori 转换回 NumPy 数组，以便进行后续处理或保存。img_src = np.asarray(img_ori)# FPS counter# 使用上一回答中提到的 CalcFPS 类的 update 方法更新帧率计算器。这里， 1.0 / (t2 - t1) 计算了上一帧图像处理的时间间隔的倒数，即当前帧的帧率，并更新到 fps_calculator 中。fps_calculator.update(1.0 / (t2 - t1))# 调用 CalcFPS 类的 accumulate 方法来计算平均帧率。avg_fps = fps_calculator.accumulate()# 检查 self.files 的类型是否为视频。这里的 self.files 可能是一个包含视频流信息的对象， type 属性表示文件类型。if self.files.type == 'video':# 如果类型为视频，则调用 draw_text 方法在视频帧上绘制文本。# def draw_text(img, text, font=cv2.FONT_HERSHEY_SIMPLEX, pos=(0, 0), font_scale=1, font_thickness=2, text_color=(0, 255, 0), text_color_bg=(0, 0, 0),):# -> 它用于在图像上绘制文本，并在其下方绘制一个背景矩形以突出显示文本。返回文本的宽度和高度。# -> return text_sizeself.draw_text(img_src,f"FPS: {avg_fps:0.1f}",pos=(20, 20),font_scale=1.0,text_color=(204, 85, 17),text_color_bg=(255, 255, 255),font_thickness=2,)# 检查 view_img 变量是否为 True ，这个变量控制是否在窗口中显示图像。if view_img:# 检查当前图像的路径 img_path 是否已经存在于 windows 列表中。 windows 列表用于跟踪已经创建的窗口，以避免重复创建。if img_path not in windows:# 如果窗口不存在，则将 img_path 添加到 windows 列表中。windows.append(img_path)# 创建一个名为 img_path 的新窗口，并设置窗口属性。 cv2.WINDOW_NORMAL 允许用户调整窗口大小， cv2.WINDOW_KEEPRATIO 保持窗口的宽高比。cv2.namedWindow(str(img_path), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux) 允许调整窗口大小（Linux）。# 设置窗口的大小为图像的宽度和高度。cv2.resizeWindow(str(img_path), img_src.shape[1], img_src.shape[0])# 在创建的窗口中显示图像 img_src 。cv2.imshow(str(img_path), img_src)# cv2.imshow(winname, img)# cv2.imshow 是 OpenCV 库中的一个函数，用于在窗口中显示图像。这个函数创建一个窗口并显示指定的图像，如果窗口不存在，则创建它；如果窗口已存在，则显示在该窗口中。# 参数 ：# winname ：窗口的名称，是一个字符串。如果名称相同，则会显示在同一个窗口中。# img ：要显示的图像，是一个 numpy 数组。# 返回值 ：# 该函数没有返回值。# 说明 ：# cv2.imshow 函数创建一个名为 winname 的窗口，并在其中显示 img 图像。# 如果窗口不存在，则创建它；如果窗口已存在，则显示在该窗口中。 # 窗口的大小默认与图像大小相同。如果需要调整窗口大小，可以使用 cv2.resizeWindow 或 cv2.namedWindow 函数设置窗口属性。# 显示图像时，OpenCV 会阻塞程序执行，直到窗口被关闭或超时。可以使用 cv2.waitKey 函数设置超时时间。# 等待1毫秒的键盘事件。这允许 OpenCV 处理窗口事件（如关闭窗口），而不会导致程序完全冻结。1毫秒的超时时间意味着程序将继续执行，而不会等待用户输入。cv2.waitKey(1) # 1 millisecond 1 毫秒。# 保存结果（带有检测的图像）。# Save results (image with detections)# 一个布尔值，指示是否需要保存图像或视频。if save_img:# 如果文件类型为 'image'，则直接使用 cv2.imwrite 将图像保存到 save_path 指定的路径。if self.files.type == 'image':# success = cv2.imwrite(filename, img, [params])# cv2.imwrite 是 OpenCV 库中的一个函数，用于将图像写入文件。这个函数以指定的文件名保存图像，支持多种图像格式，如 BMP、JPEG、PNG 等。# 参数 ：# filename ：要写入的图像文件的名称，包括路径。# img ：要保存的图像，是一个 numpy 数组。# params ：一个可选参数列表，用于指定图像编码的额外参数。例如，对于 JPEG 图像，可以指定质量参数。# 返回值 ：# success ：一个布尔值，如果图像成功写入文件，则返回 True ；否则返回 False 。# 说明 ：# cv2.imwrite 函数将 img 图像保存到 filename 指定的文件中。如果文件已存在，则会被覆盖。# 支持的图像格式包括 BMP、JPG、PNG、TIFF 等，具体支持的格式取决于 OpenCV 编译时的配置。# 通过指定 params 参数，可以控制图像的编码质量。例如，对于 JPEG 图像，可以指定压缩质量。cv2.imwrite(save_path, img_src)# 如果文件类型为 'video' 或 'stream'，则执行以下操作。else: # 'video' or 'stream' “视频”或“流”。# 如果 save_path 发生变化，表示需要开始一个新的视频文件。if vid_path != save_path: # new video 新视频。# 更新 vid_path 为新的保存路径。vid_path = save_path# 如果 vid_writer 已经存在，释放之前的 VideoWriter 对象。if isinstance(vid_writer, cv2.VideoWriter):# 释放之前的 vid_writer 。vid_writer.release() # release previous video writer# 如果 vid_cap 存在，即处理的是视频文件。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))# 如果 vid_cap 不存在，即处理的是视频流。else: # stream 流。# 设置默认的帧率、宽度和高度。fps, w, h = 30, img_ori.shape[1], img_ori.shape[0]# 确保保存路径的文件扩展名为 '.mp4'。save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos 强制结果视频使用 *.mp4 后缀。# 创建一个新的 cv2.VideoWriter 对象，用于写入视频。vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))# 将图像 img_src 作为视频帧写入视频文件。vid_writer.write(img_src)# 它用于在图像推理之前对图像进行预处理。@staticmethod# 1.img_src ：原始图像数据，可以是图像文件路径或图像数组。# 2.img_size ：目标图像尺寸，用于调整图像大小。# 3.stride ：模型的步长，用于确定图像调整时的对齐。# 4.half ：布尔值，指示是否将图像转换为半精度（FP16）。def process_image(img_src, img_size, stride, half):# 在图像推理之前处理图像。'''Process image before image inference.'''# def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):# -> 在满足步幅倍数约束的同时调整图像大小并填充。# -> return im, r, (left, top)# 使用 letterbox 函数将图像调整到 img_size 指定的尺寸，同时保持图像的长宽比。 letterbox 函数返回一个包含调整后的图像和一些额外信息的元组，这里只取第一个元素，即调整后的图像。image = letterbox(img_src, img_size, stride=stride)[0]# Convert# 将图像从 HWC（高度、宽度、通道）格式转换为 CHW（通道、高度、宽度）格式，并将 BGR 颜色空间转换为 RGB。image = image.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB# np.ascontiguousarray(a, dtype=None)# 是 NumPy 库中的一个函数，用于返回一个连续存储的数组（在内存中是按行优先顺序存储的）。# 如果输入数组已经是连续的，它将返回输入数组的引用；否则，它会返回一个新的、连续的副本。# a ：要转换的输入数组。# dtype（可选） ：所需的数组数据类型。如果未指定，则保持输入数组的数据类型。# 返回 ：返回一个内存中连续存储的数组。# 将 NumPy 数组转换为 PyTorch 张量，并确保数组是连续的。image = torch.from_numpy(np.ascontiguousarray(image))# 如果 half 为 True ，则将图像转换为半精度（FP16），否则转换为全精度（FP32）。image = image.half() if half else image.float() # uint8 to fp16/32# 将图像的像素值从 [0, 255] 范围归一化到 [0.0, 1.0] 范围。image /= 255 # 0 - 255 to 0.0 - 1.0# 返回值 ：# 1.image ：预处理后的图像张量。# 2.img_src ：原始图像数据或路径。return image, img_src# 它用于将检测框的坐标从模型输入图像的尺寸重新缩放到原始图像的尺寸。@staticmethod# ori_shape ：原始图像的尺寸，格式为 (高度, 宽度) 。# boxes ：检测框的坐标，格式为 (x1, y1, x2, y2) ，其中 (x1, y1) 是左上角坐标， (x2, y2) 是右下角坐标。# target_shape ：模型输入图像的尺寸，格式为 (高度, 宽度) 。def rescale(ori_shape, boxes, target_shape):# 将输出重新缩放为原始图像形状。'''Rescale the output to the original image shape'''# 计算原始图像和目标图像之间的最小缩放比例。ratio = min(ori_shape[0] / target_shape[0], ori_shape[1] / target_shape[1])# 计算在宽度和高度方向上的填充量，以便将目标图像居中放置在原始图像中。padding = (ori_shape[1] - target_shape[1] * ratio) / 2, (ori_shape[0] - target_shape[0] * ratio) / 2# 将检测框的 x 坐标向左移动填充量。boxes[:, [0, 2]] -= padding[0]# 将检测框的 y 坐标向上移动填充量。boxes[:, [1, 3]] -= padding[1]# 将检测框的坐标按缩放比例进行缩放。boxes[:, :4] /= ratio# 确保检测框的左上角 x 坐标在目标图像宽度范围内。boxes[:, 0].clamp_(0, target_shape[1]) # x1# 确保检测框的左上角 y 坐标在目标图像高度范围内。boxes[:, 1].clamp_(0, target_shape[0]) # y1# 确保检测框的右下角 x 坐标在目标图像宽度范围内。boxes[:, 2].clamp_(0, target_shape[1]) # x2# 确保检测框的右下角 y 坐标在目标图像高度范围内。boxes[:, 3].clamp_(0, target_shape[0]) # y2# 重新缩放后的检测框坐标。return boxes# 作用是确保输入的图像尺寸 img_size 是模型步长 s 的倍数。这个功能在深度学习模型中很常见，因为许多模型要求输入图像的尺寸必须是特定的步长（例如32、64等）的倍数。# 1.self ：类的实例本身。# 2.img_size ：输入的图像尺寸，可以是一个整数或一个列表。# 3.s ：步长，默认为32。# 4.floor ：最小尺寸限制，默认为0。def check_img_size(self, img_size, s=32, floor=0):# 确保图像大小在每个维度上都是 stride 的倍数，并返回一个新的图像形状列表。"""Make sure image size is a multiple of stride s in each dimension, and return a new shape list of image."""# 如果 img_size 是一个整数（例如640），则表示图像的宽度和高度相同。# 使用 self.make_divisible 方法（这个方法未在代码中定义，但应该是一个自定义方法）来确保每个维度的尺寸是步长 s 的倍数。# max 函数确保新的尺寸不会低于 floor 指定的最小尺寸。if isinstance(img_size, int): # integer i.e. img_size=640 整数，即 img_size=640。new_size = max(self.make_divisible(img_size, int(s)), floor)# 如果 img_size 是一个列表（例如[640, 480]），则表示图像的宽度和高度不同。elif isinstance(img_size, list): # list i.e. img_size=[640, 480] 列表，即 img_size=[640, 480]。new_size = [max(self.make_divisible(x, int(s)), floor) for x in img_size]else:# 如果 img_size 既不是整数也不是列表，抛出异常。raise Exception(f"Unsupported type of img_size: {type(img_size)}") # 不支持的 img_size 类型：{type(img_size)}。if new_size != img_size:# 如果原始尺寸和计算后的新尺寸不同，打印一条警告信息，提示用户图像尺寸已被调整。print(f'WARNING: --img-size {img_size} must be multiple of max stride {s}, updating to {new_size}') # 警告：--img-size {img_size} 必须是最大步幅 {s} 的倍数，更新为 {new_size}# 如果 img_size 是列表，直接返回调整后的尺寸列表。# 如果 img_size 是整数，返回一个包含两个相同值的列表，这两个值都是调整后的尺寸。return new_size if isinstance(img_size,list) else [new_size]*2def make_divisible(self, x, divisor):# 向上修正 x 的值，使其可以被除数整除。# Upward revision the value x to make it evenly divisible by the divisor.return math.ceil(x / divisor) * divisor# 它用于在图像上绘制文本，并在其下方绘制一个背景矩形以突出显示文本。@staticmethod# img ：要绘制文本的原始图像。# text ：要绘制的文本字符串。# font ：用于绘制文本的字体，默认为 cv2.FONT_HERSHEY_SIMPLEX 。# pos ：文本在图像上的位置，默认为 (0, 0) 。# font_scale ：字体的缩放比例，默认为 1 。# font_thickness ：字体的线条宽度，默认为 2 。# text_color ：文本的颜色，默认为 (0, 255, 0) ，即绿色。# text_color_bg ：文本背景的颜色，默认为 (0, 0, 0) ，即黑色。def draw_text(img,text,font=cv2.FONT_HERSHEY_SIMPLEX,pos=(0, 0),font_scale=1,font_thickness=2,text_color=(0, 255, 0),text_color_bg=(0, 0, 0),):# 设置文本背景的偏移量。offset = (5, 5)# 获取文本的位置坐标。x, y = pos# (cv2.Size(width, height), baseline) = cv2.getTextSize(text, fontFace, fontScale, thickness)# cv2.getTextSize 函数是 OpenCV 库中的一个函数，它用于计算给定文本的尺寸（宽度和高度）。这个函数在绘制文本之前非常有用，因为它可以帮助你确定文本的尺寸，从而可以进行适当的定位和布局。# 参数 ：# text ：要测量的文本字符串。# fontFace ：字体类型，可以是 OpenCV 预定义的字体，如 cv2.FONT_HERSHEY_SIMPLEX 、 cv2.FONT_HERSHEY_PLAIN 等。# fontScale ：字体缩放因子，用于调整字体大小。# thickness ：字体线条的厚度。# 返回值 ：# cv2.Size(width, height) ：一个 cv2.Size 对象，包含文本的宽度和高度。# baseline ：文本基线（即文本底部到基线的距离），这个值可以用来确定文本的位置，以确保基线对齐。# 使用 OpenCV 的 getTextSize 函数计算文本的宽度和高度。text_size, _ = cv2.getTextSize(text, font, font_scale, font_thickness)# 获取文本的宽度和高度。text_w, text_h = text_size# 计算背景矩形的左上角坐标。rec_start = tuple(x - y for x, y in zip(pos, offset))# 计算背景矩形的右下角坐标。rec_end = tuple(x + y for x, y in zip((x + text_w, y + text_h), offset))# 使用 OpenCV 的 rectangle 函数绘制填充的背景矩形。cv2.rectangle(img, rec_start, rec_end, text_color_bg, -1)# 使用 OpenCV 的 putText 函数在图像上绘制文本。cv2.putText(img,text,# 计算文本的基线位置，确保文本位于背景矩形的上方。(x, int(y + text_h + font_scale - 1)),font,font_scale,text_color,font_thickness,cv2.LINE_AA,)# 返回文本的宽度和高度。return text_size# 它用于在图像上绘制一个边界框（bounding box）以及与之关联的标签（label）。@staticmethod# 1.image ：要绘制边界框和标签的原始图像。# 2.lw ：边界框的线条宽度。# 3.box ：边界框的坐标，格式为 (x1, y1, x2, y2) 。# 4.label ：边界框内的标签文本，默认为空字符串。# 5.color ：边界框的颜色，默认为 (128, 128, 128) ，即灰色。# 6.txt_color ：标签文本的颜色，默认为 (255, 255, 255) ，即白色。# 7.font ：用于绘制文本的字体，默认为 cv2.FONT_HERSHEY_COMPLEX 。def plot_box_and_label(image, lw, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255), font=cv2.FONT_HERSHEY_COMPLEX):# 将一个 xyxy 框添加到带标签的图像中。# Add one xyxy box to image with label# 将边界框坐标转换为整数，并分别表示左上角和右下角。p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))# 使用 OpenCV 的 rectangle 函数在图像上绘制边界框。cv2.rectangle(image, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)# 如果标签不为空，则继续绘制标签。if label:# 计算文本的线条宽度。tf = max(lw - 1, 1) # font thickness 字体粗细。# 使用 OpenCV 的 getTextSize 函数计算标签文本的宽度和高度。w, h = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0] # text width, height 文字宽度、高度。# 判断标签是否适合放在边界框的上方。outside = p1[1] - h - 3 >= 0 # label fits outside box 标签适合边界框外面。# 根据标签位置计算背景矩形的右下角坐标。p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3# 使用 OpenCV 的 rectangle 函数绘制填充的背景矩形。cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA) # filled 填充。# 使用 OpenCV 的 putText 函数在图像上绘制标签文本。cv2.putText(image, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), font, lw / 3, txt_color,thickness=tf, lineType=cv2.LINE_AA)# 该方法没有返回值，它直接在输入的 image 上进行绘制。# 用于检查指定路径的字体文件是否存在，并返回一个 PIL（Python Imaging Library）TrueType 字体对象。如果字体文件不存在，它将尝试下载字体并加载。@staticmethod# font ：字体文件的路径，默认为 './yolov6/utils/Arial.ttf' 。# size ：字体大小，默认为 10 。def font_check(font='./yolov6/utils/Arial.ttf', size=10):# 返回 PIL TrueType 字体，必要时下载到 CONFIG_DIR 。# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary# 使用 osp.exists 函数（ osp 是 os.path 的别名）检查字体文件是否存在。如果文件不存在，将抛出一个 AssertionError 。assert osp.exists(font), f'font path not exists: {font}' # 字体路径不存在：{font}。# 尝试执行以下代码块。try:# 使用 PIL 的 ImageFont.truetype 方法加载字体文件，并返回一个字体对象。这里有一个逻辑错误，因为 font 是一个路径字符串，而不是一个具有 exists 方法的对象，所以 font.exists() 将导致错误。return ImageFont.truetype(str(font) if font.exists() else font.name, size)# 如果加载字体时发生任何异常，将捕获异常。except Exception as e: # download if missing# 在异常处理中再次尝试加载字体。这里假设即使在异常处理中，也不需要下载字体，而是直接加载。return ImageFont.truetype(str(font), size)# 它用于将检测框的坐标从 [x1, y1, x2, y2] 格式（其中 (x1, y1) 是左上角坐标， (x2, y2) 是右下角坐标）转换为 [x, y, w, h] 格式（其中 (x, y) 是中心点坐标， w 是宽度， h 是高度）。@staticmethod# x ：检测框的坐标，格式为 [n, 4] ，其中 n 是检测框的数量。def box_convert(x):# 将形状为 [n, 4] 的框从 [x1, y1, x2, y2] 转换为 [x, y, w, h]，其中 x1y1=左上角，x2y2=右下角# Convert boxes with shape [n, 4] from [x1, y1, x2, y2] to [x, y, w, h] where x1y1=top-left, x2y2=bottom-right# 如果输入 x 是 PyTorch 张量，则使用 clone() 方法创建一个副本；如果 x 是 NumPy 数组，则使用 np.copy() 创建一个副本。这一步确保了原始输入数据不会被修改。y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)# 计算 x 坐标的中心点。y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center# 计算 y 坐标的中心点。y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center# 计算检测框的宽度。y[:, 2] = x[:, 2] - x[:, 0] # width# 计算检测框的高度。y[:, 3] = x[:, 3] - x[:, 1] # height# y ：转换后的检测框坐标，格式为 [x, y, w, h] 。return y# 它用于生成一组预定义颜色的调色板，并根据给定的索引 i 返回相应的颜色。@staticmethod# 1.i ：颜色索引，用于从调色板中选择颜色。# 2.bgr ：布尔值，指示返回的颜色格式是否为 BGR（Blue Green Red），默认为 False，即返回 RGB（Red Green Blue）格式。def generate_colors(i, bgr=False):# hex ：包含多个十六进制颜色代码的元组。hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB','2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')# 初始化一个空列表来存储 RGB 颜色元组。palette = []# 遍历 hex 中的每个十六进制颜色代码。for iter in hex:# 在十六进制颜色代码前加上 # 。h = '#' + iter# 将十六进制颜色代码转换为 RGB 元组，并添加到 palette 列表中。palette.append(tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)))# 获取调色板中颜色的数量。num = len(palette)# 根据索引 i 从调色板中选择颜色，使用模运算确保索引在有效范围内。color = palette[int(i) % num]# 如果 bgr 为 True，则返回 BGR 格式的颜色元组；否则返回 RGB 格式的颜色元组。# 根据 bgr 参数返回 RGB 或 BGR 格式的颜色元组。return (color[2], color[1], color[0]) if bgr else color3.class CalcFPS:
# 定义了一个名为 CalcFPS 的类，它用于计算帧率（Frames Per Second）。这个类可以帮助你测量和跟踪处理视频流或图像序列时的性能。
class CalcFPS:def __init__(self, nsamples: int = 50):# 构造函数，初始化一个 deque （双端队列）来存储最近 nsamples 个样本的帧处理时间。# nsamples ：队列的最大长度，默认为50。self.framerate = deque(maxlen=nsamples)# 更新方法，用于添加新的帧处理时间到队列中。# duration ：处理一帧图像所花费的时间，以秒为单位。def update(self, duration: float):self.framerate.append(duration)# 计算并返回平均帧处理时间。# 如果队列中有超过一个样本，使用 np.average 计算平均值。# 如果队列中只有一个样本或为空，返回0.0。def accumulate(self):if len(self.framerate) > 1:return np.average(self.framerate)else:return 0.0
