【YOLOv3】源码（detect.py）

最新推荐文章于 2025-01-05 12:16:35 发布

gma999

最新推荐文章于 2025-01-05 12:16:35 发布

阅读量683

点赞数 34

文章标签： YOLO

本文链接：https://blog.youkuaiyun.com/gma999/article/details/144652298

版权

概述

该文件类似于实时监控与检测团队，负责在建筑项目中实时监控施工进展、检测潜在问题并记录和报告结果

主要模块

参数解析与初始化
- 功能：解析命令行参数，设置检测配置
- 实时监控团队制定详细的监控计划和资源配置
模型加载与配置
- 功能：加载模型和配置文件，配置设备，设置检测工具参数
- 实时监控团队选择和配置检测工具，确保检测过程高效和准确
数据加载与源解析
- 功能：根据输入源加载数据，配置数据加载器，进行设备优化和预热
- 实时监控团队选择不同的监控源，配置相应的检测流程，确保数据的顺利输入和处理
推理循环
- 功能：对每张图像或视频帧进行推理，应用NMS，处理和记录检测结果，保存和可视化检测结果
- 实时监控团队对每个监控点或区域进行实际检测和记录，确保施工过程中的每个部分符合设计标准
结果处理与保存
- 功能：保存检测结果到文本文件，绘制和保存边界框，裁剪并保存检测到的目标区域
- 实时监控团队记录检测到的问题，标记问题位置，并根据需要保存详细的检测图片，便于后续分析和修复

主要模块

参数解析与初始化

与train.py文件类似，详细参考代码注释

def parse_opt():
    """
    解析命令行参数，设置检测相关的所有参数。
    """
    parser = argparse.ArgumentParser()

    # ----------------------------- 模型相关参数 -----------------------------------
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'weight/yolov3.pt',
                        help='模型路径，可以是单个路径或多个路径列表')  # 权重文件路径
    parser.add_argument('--source', type=str, default=ROOT / 'data/images',
                        help='输入数据来源：可以是文件、文件夹、URL 或摄像头 (0 表示摄像头)')  # 输入数据路径
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[416],
                        help='推理时的图片尺寸（可以是单个数值，也可以是 [h, w] 格式）')  # 图片输入尺寸

    # ----------------------------- 推理相关参数 -----------------------------------
    parser.add_argument('--conf-thres', type=float, default=0.6, help='置信度阈值')  # 检测框置信度阈值
    parser.add_argument('--iou-thres', type=float, default=0.5, help='非极大值抑制（NMS）的 IoU 阈值')  # NMS 阈值
    parser.add_argument('--max-det', type=int, default=1000, help='每张图像的最大检测数量')  # 最大检测框数量
    parser.add_argument('--device', default='', help='指定设备，例如 "0" 或 "0,1,2,3" 或 "cpu"')  # 设备选择

    # ----------------------------- 结果相关参数 -----------------------------------
    parser.add_argument('--view-img', action='store_true', default=True, help='显示检测结果')  # 是否显示结果
    parser.add_argument('--save-txt', action='store_true', default=True, help='将检测结果保存为 *.txt 文件')  # 保存结果为文本
    parser.add_argument('--save-conf', action='store_true', default=True, help='保存检测结果的置信度值')  # 保存置信度
    parser.add_argument('--save-crop', action='store_true', default=True, help='保存裁剪后的预测框图像')  # 保存裁剪结果
    parser.add_argument('--nosave', action='store_true', help='不保存检测的图片或视频')  # 禁止保存图片/视频
    parser.add_argument('--classes', nargs='+', type=int,
                        help='根据类别进行过滤：--classes 0 或 --classes 0 2 3')  # 只检测指定的类别
    parser.add_argument('--agnostic-nms', action='store_true', help='使用类别无关的 NMS')  # 类别无关的 NMS
    parser.add_argument('--augment', action='store_true', help='使用增强推理')  # 推理时是否使用数据增强
    parser.add_argument('--visualize', action='store_true', help='可视化特征')  # 特征可视化

    # ----------------------------- 其他配置参数 -----------------------------------
    parser.add_argument('--update', action='store_true', help='更新所有模型')  # 更新所有模型（仅适用于特定流程）
    parser.add_argument('--project', default=ROOT / 'runs/detect', help='保存结果的路径')  # 检测结果保存路径
    parser.add_argument('--name', default='exp', help='保存结果的子目录名')  # 检测结果保存子目录
    parser.add_argument('--exist-ok', action='store_true', default=False,
                        help='如果目录已存在，不创建新目录')  # 不创建新目录，直接使用现有目录
    parser.add_argument('--line-thickness', default=3, type=int, help='边界框的线条粗细（像素）')  # 检测框的线条粗细
    parser.add_argument('--hide-labels', default=False, action='store_true', help='隐藏检测结果的类别标签')  # 隐藏类别标签
    parser.add_argument('--hide-conf', default=False, action='store_true', help='隐藏检测结果的置信度值')  # 隐藏置信度值

    # ----------------------------- 高级参数 -----------------------------------
    parser.add_argument('--half', action='store_true', help='使用 FP16 半精度推理')  # 是否使用半精度推理
    parser.add_argument('--dnn', action='store_true', help='使用 OpenCV 的 DNN 模块进行 ONNX 推理')  # 是否使用 OpenCV 推理

    # 解析参数
    opt = parser.parse_args()

    # 如果输入图像尺寸只有一个数值，则自动扩展为 [h, w] 格式
    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1

    # 打印解析后的参数
    print_args(FILE.stem, opt)

    return opt

模型加载与配置

主要功能则是加载模型权重和配置文件，设置模型参数，配置设备（CPU/GPU），初始化推理数据。类似于实时监控与检测团队选择和配置检测工具、准备检测资料

位于YOLOv3流程中的前置准备阶段

主要流程总结

模型加载：加载 YOLOv3 模型的权重、设置设备和精度
数据加载：根据输入源类型（摄像头或文件）加载数据
模型预热：在开始推理之前运行一次“空推理”以预热模型，主要是为了让后面的推理速度加快

# 载入模型和模型参数并调整模型 

# 加载模型
device = select_device(device)  # 选择设备，支持 GPU 或 CPU
model = DetectMultiBackend(weights, device=device, dnn=dnn)  # 加载模型，支持多种后端（PyTorch、ONNX 等）
stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx  # 获取模型的相关属性
# stride: 模型的步长，用于调整输入图像尺寸；names: 类别名称列表；
# pt: 是否是 PyTorch 模型；jit: 是否是 JIT 模型；onnx: 是否是 ONNX 模型
imgsz = check_img_size(imgsz, s=stride)  # 检查并调整输入图像尺寸，使其符合模型的步长要求

# 设置半精度推理
half &= pt and device.type != 'cpu'  # 半精度仅适用于 CUDA 上的 PyTorch 模型
if pt:  # 如果模型是 PyTorch 模型
    model.model.half() if half else model.model.float()  # 将模型设置为半精度 (float16) 或全精度 (float32)

数据加载与源解析

根据输入源加载相应的数据，配置数据加载器，从而确保检测过程中数据的顺利输入和处理。这类似于实时监控与检测团队选择和配置不同的监控源，如实时摄像头、预先录制的视频或图片资料

# ===================================== 3、加载推理数据 =====================================

# 设置数据加载器
if webcam:  # 如果输入源是摄像头
    view_img = check_imshow()  # 检查当前环境是否支持图像显示（例如 GUI 环境是否可用）
    cudnn.benchmark = True  # 启用 cuDNN 的 benchmark 模式
    # 作用：当输入数据尺寸固定时，启用 benchmark 可以提升推理速度
    dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)
    # `LoadStreams` 用于加载实时流数据（例如摄像头输入或视频流）。
    # 参数：
    # - `source`：数据源，可以是摄像头或视频流。
    # - `img_size`：输入图像大小。
    # - `stride`：模型步长，用于调整输入图像尺寸。
    # - `auto`：是否自动调整输入图像尺寸以匹配模型要求。
    bs = len(dataset)  # 设置批量大小为流的数量（多个摄像头或视频流的数量）
else:  # 如果输入源是文件（例如图片或视频文件）
    dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
    # `LoadImages` 用于从本地文件加载数据（图片或视频文件）。
    # 参数与 `LoadStreams` 类似。
    bs = 1  # 设置批量大小为 1，因为处理文件时通常是逐帧进行

# 初始化视频路径和写入器
vid_path, vid_writer = [None] * bs, [None] * bs
# `vid_path`：存储视频文件的路径，用于保存处理后的推理结果。
# `vid_writer`：视频写入器，用于将推理结果写入视频文件。

# 如果使用 PyTorch 模型（pt=True）且设备类型不是 CPU：
if pt and device.type != 'cpu':
    # 模型预热
    model(
        torch.zeros(1, 3, *imgsz)  # 创建一个全零张量，形状为 (1, 3, img_height, img_width)
        .to(device)  # 将张量移动到指定设备（GPU 或 CPU）
        .type_as(next(model.model.parameters()))  # 确保张量的数据类型与模型参数一致（float16 或 float32）
    )  # 运行一次模型推理（空推理），预热模型
    # 预热的目的是加载必要的计算资源，并优化后续推理的性能

# 初始化变量
dt, seen = [0.0, 0.0, 0.0], 0
# `dt`：记录时间消耗的列表，其中包含三个阶段的时间：
# - dt[0]：数据加载时间
# - dt[1]：推理时间
# - dt[2]：后处理（如 NMS）时间
# `seen`：已处理的图像数量，用于统计推理进度

推理循环

主要步骤分析

图像预处理：对待检测的部分进行标准化和预处理，从而确保检测工具可以准确的识别和分析
预测：开始实际的检测
非极大值抑制：过滤重复的检测结果，保证每一个区域只检测一次
结果处理与保存
可视化显示
结果输出

# ===================================== 5、正式推理 =====================================

# 遍历数据集，逐帧或逐图片处理
for path, im, im0s, vid_cap, s in dataset:
    """
    path: 当前帧或图片的路径
    im: 调整尺寸并添加填充后的图片（用于推理的输入）
    im0s: 原始图片（未经过调整和填充）
    vid_cap: 如果是视频文件，为视频捕获对象；如果是图片，则为 None
    s: 输出的初始字符串信息
    """

    # 5.1 获取当前时间，用于计算处理时间
    t1 = time_sync()

    # 5.2 将图片转换为 PyTorch 张量并加载到指定设备
    im = torch.from_numpy(im).to(device)
    im = im.half() if half else im.float()  # 如果启用半精度，转换为 fp16，否则为 fp32
    im /= 255  # 归一化像素值到 [0, 1]
    if len(im.shape) == 3:  # 如果图片是 3 维（无 batch 维度）
        im = im[None]  # 添加 batch 维度，变成 [1, channels, height, width]
    t2 = time_sync()  # 获取当前时间
    dt[0] += t2 - t1  # 累计数据预处理时间

    # 如果启用了可视化，将设置保存路径并创建目录
    visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False

    # 5.3 使用模型进行前向推理
    pred = model(im, augment=augment, visualize=visualize)  # 推理预测
    t3 = time_sync()
    dt[1] += t3 - t2  # 累计推理时间

    # 5.4 非极大值抑制 (NMS)
    pred = non_max_suppression(
        pred,
        conf_thres,  # 置信度阈值
        iou_thres,  # IoU 阈值
        classes,  # 是否仅保留特定类别
        agnostic_nms,  # 是否进行类别无关的 NMS
        max_det=max_det  # 每张图片的最大检测目标数
    )
    dt[2] += time_sync() - t3  # 累计 NMS 时间

    # 5.5 第二阶段分类器（可选）
    # 如果启用了第二阶段分类器，则将结果传递给分类器进行进一步处理
    # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

    # 5.6 后续保存或打印预测信息
    # 遍历每张图片的预测结果
    for i, det in enumerate(pred):  # 对于每张图片的预测结果
        seen += 1  # 已处理的图片数
        if webcam:  # 如果数据来源是摄像头或实时流
            p, im0, frame = path[i], im0s[i].copy(), dataset.count
            s += f'{i}: '
        else:  # 如果数据来源是图片或视频文件
            p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

        # 转换路径格式
        p = Path(p)  # 将路径字符串转换为 Path 对象
        save_path = str(save_dir / p.name)  # 保存图片的路径
        txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')
        s += '%gx%g ' % im.shape[2:]  # 输出推理的图片尺寸
        gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # 归一化比例，用于将检测框还原到原始图片尺寸
        imc = im0.copy() if save_crop else im0  # 用于保存裁剪目标的图片
        annotator = Annotator(im0, line_width=line_thickness, example=str(names))  # 用于绘制检测框的工具

        if len(det):  # 如果有检测结果
            # 将检测框从 img_size 映射回原图尺寸
            det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

            # 打印每个类别的检测结果
            for c in det[:, -1].unique():
                n = (det[:, -1] == c).sum()  # 每个类别的检测数量
                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # 添加检测信息到字符串

            # 保存检测结果到文件或图像
            for *xyxy, conf, cls in reversed(det):  # 遍历每个检测框
                if save_txt:  # 保存到文本文件
                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # 转换为归一化 xywh
                    line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # 根据配置添加置信度
                    with open(txt_path + '.txt', 'a') as f:
                        f.write(('%g ' * len(line)).rstrip() % line + '\n')

                # 如果需要保存裁剪目标或绘制检测框
                if save_img or save_crop or view_img:
                    c = int(cls)  # 类别索引
                    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))  # 绘制检测框
                    if save_crop:  # 保存裁剪后的目标
                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

        # 打印处理时间
        print(f'{s}Done. ({t3 - t2:.3f}s)')

        # 如果启用了查看功能，则显示检测结果
        im0 = annotator.result()
        if view_img:  # 显示检测结果
            cv2.imshow(str(p), im0)
            cv2.waitKey(1)

        # 如果需要保存图片或视频结果
        if save_img:
            if dataset.mode == 'image':  # 如果是图片
                cv2.imwrite(save_path, im0)  # 保存检测后的图片
            else:  # 如果是视频
                if vid_path[i] != save_path:  # 如果路径发生变化，重新创建视频写入器
                    vid_path[i] = save_path
                    if isinstance(vid_writer[i], cv2.VideoWriter):
                        vid_writer[i].release()  # 释放之前的视频写入器
                    if vid_cap:  # 视频文件
                        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:  # 实时流
                        fps, w, h = 30, im0.shape[1], im0.shape[0]
                        save_path += '.mp4'
                    vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                vid_writer[i].write(im0)  # 写入当前帧

# 打印推理速度
t = tuple(x / seen * 1E3 for x in dt)  # 计算每张图片的平均处理时间（毫秒）
print(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_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}"
         if save_txt else '')
    LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")

# 如果启用了更新选项，则移除优化器信息
if update:
    strip_optimizer(weights)  # 去除优化器信息，减小模型大小

结果处理与保存

将检测到的目标结果保存为文本文件（可选），在图像上绘制边界框和标签，并根据需要保存裁剪后的目标图像

类似于实时监控与检测团队记录检测到的问题、标记问题位置，并根据需要保存问题区域的详细图片

总结

保存预测信息：将检测到的问题记录在报告中，包含问题类别、位置、置信度等信息
绘制边界框：在原始图像或视频帧上标记检测到的问题区域，类似于在施工图上标注发现的质量问题
裁剪并保存问题区域：将检测到的问题区域剪切出来保存为独立的图片，便于后续的详细分析和修复

for *xyxy, conf, cls in reversed(det):  # 遍历每个检测结果（从后往前）
    """
    xyxy: 检测框的左上角和右下角坐标 (x1, y1, x2, y2)
    conf: 检测框的置信度分数
    cls: 检测框的类别索引
    det: 包含所有检测结果的张量，形状为 [num_detections, 6]，每行是 [x1, y1, x2, y2, conf, cls]
    """

    # 保存预测信息到文本文件
    if save_txt:  # 如果启用了保存文本选项
        # 将检测框坐标 (xyxy) 转换为中心点坐标和宽高格式 (xywh)，并进行归一化
        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  
        """
        xyxy2xywh：将边界框坐标从 (x1, y1, x2, y2) 格式转换为 (x_center, y_center, width, height)
        gn：归一化因子，是 [width, height, width, height]，用于将绝对坐标归一化到 0-1 之间
        .view(1, 4)：将 xyxy 转换为 1 行 4 列的张量
        .tolist()：将张量转换为 Python 列表
        """

        # 根据是否需要保存置信度，生成保存行的格式
        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # 保存类别、归一化后的坐标和置信度

        # 将检测结果写入对应的文本文件
        with open(txt_path + '.txt', 'a') as f:  # 打开或创建 .txt 文件
            f.write(('%g ' * len(line)).rstrip() % line + '\n')  # 按格式写入每一行

    # 在图像上绘制检测框或保存裁剪的目标区域
    if save_img or save_crop or view_img:  # 如果需要保存图片、裁剪目标或显示结果
        c = int(cls)  # 将类别索引转换为整数
        # 根据配置，生成标签内容（类别名称，或类别名称加置信度）
        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))  
        """
        annotator.box_label：在图像上绘制边界框并添加类别标签
        xyxy：边界框的左上角和右下角坐标
        label：要显示的标签
        colors(c, True)：为当前类别选择颜色
        """

        # 如果启用了裁剪保存选项
        if save_crop:
            # 裁剪检测框对应的目标区域，并保存为图片
            save_one_box(
                xyxy,  # 边界框坐标
                imc,  # 原始图像的副本
                file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg',  # 保存路径
                BGR=True  # 保存为 BGR 格式（适用于 OpenCV）
            )
            """
            save_one_box：
            - 裁剪目标区域并保存为单独的图片
            - file：保存路径，包含类别名称文件夹和图片文件名
            - BGR：裁剪的图片保存为 BGR 格式，适配 OpenCV 的格式需求
            """