目标检测经典模型之YOLOV5-detect.py源码解析（持续更新）

MickeyCV2024-07-18 12:11

detect文件框架

本帖是YOLOV5推理部分代码的中文逐行注释。由于AI注释的缘故，可能与源码会有小部分出入，所以不建议复制粘贴替换源码的detect.py文件。本贴的初衷是YOLOV5源码逻辑的学习，后续会不断修正该代码和加入新的注释。

一、导入模块包

python 复制代码 
import argparse  # 引入argparse库，用于解析命令行参数
import os  # 引入os库，用于进行操作系统相关的操作
import platform  # 引入platform库，用于获取平台信息
import sys  # 引入sys库，用于操作Python运行时环境
from pathlib import Path  # 引入Path库，用于处理文件和目录路径

import cv2  # 引入OpenCV库，用于图像处理
import torch  # 引入PyTorch库，用于深度学习
FILE = Path(__file__).resolve()  # 获取当前文件的绝对路径
ROOT = FILE.parents[0]  # 获取当前文件的父目录路径

if str(ROOT) not in sys.path:  # 如果当前文件的父目录路径不在系统路径中
    sys.path.append(str(ROOT))  # 将当前文件的父目录路径添加到系统路径中

ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # 将ROOT路径相对于当前工作目录进行转换

from ultralytics.utils.plotting import Annotator, colors, save_one_box
from models.common import DetectMultiBackend  # 从models.common模块中导入DetectMultiBackend类
from utils.dataloaders import LoadImages, LoadStreams  # 从utils.dataloaders模块中导入LoadImages和LoadStreams类
from utils.general import (  # 从utils.general模块中导入多个函数和类
    LOGGER,
    Profile,
    check_file,
    check_img_size,
    check_imshow,
    check_requirements,
    colorstr,
    cv2,
    increment_path,
    non_max_suppression,
    print_args,
    scale_boxes,
    strip_optimizer,
    xyxy2xywh,
)
from utils.torch_utils import select_device, smart_inference_mode # 从utils.torch_utils模块中导入select_device和time_sync和smart_inference_mode函数

二、定义run函数

python 复制代码 
@smart_inference_mode()
def run(
        weights='yolov5s.pt',  # 模型权重文件路径，默认值为'yolov5s.pt'
        source='data/images',  # 输入源，可以是文件、目录、URL或摄像头，默认值为'data/images'
        data='data/coco128.yaml',  # 数据集配置文件路径，默认值为'data/coco128.yaml'
        imgsz=640,  # 输入图像尺寸，默认值为640
        conf_thres=0.25,  # 置信度阈值，默认值为0.25
        iou_thres=0.45,  # 非极大值抑制的IoU阈值，默认值为0.45
        max_det=1000,  # 每张图像的最大检测数量，默认值为1000
        device='',  # 使用的设备，可以是'cpu'或'cuda:0'，默认值为''
        view_img=False,  # 是否显示检测结果，默认值为False
        save_txt=False,  # 是否将检测结果保存到文本文件，默认值为False
        save_conf=False,  # 是否保存置信度，默认值为False
        save_crop=False,  # 是否保存裁剪后的检测框，默认值为False
        nosave=False,  # 是否保存图像/视频，默认值为False
        classes=None,  # 按类别过滤，例如0或0 2 3，默认值为None
        agnostic_nms=False,  # 是否使用类别无关的NMS，默认值为False
        augment=False,  # 是否使用增强推理，默认值为False
        visualize=False,  # 是否可视化特征，默认值为False
        update=False,  # 是否更新所有模型，默认值为False
        project='runs/detect',  # 保存结果的项目路径，默认值为'runs/detect'
        name='exp',  # 保存结果的文件夹名称，默认值为'exp'
        exist_ok=False,  # 是否允许现有项目名称，默认值为False
        line_thickness=3,  # 边界框的厚度（像素），默认值为3
        hide_labels=False,  # 是否隐藏标签，默认值为False
        hide_conf=False,  # 是否隐藏置信度，默认值为False
        half=False,  # 是否使用FP16半精度推理，默认值为False
        dnn=False,  # 是否使用OpenCV DNN进行ONNX推理，默认值为False
        vid_stride=1,  # 视频帧率步幅，默认值为1
):
	# 将source变量转换为字符串
	source = str(source)
	
	# 判断是否需要保存推理后的图像，除非指定了--nosave或source是文本文件
	save_img = not nosave and not source.endswith(".txt")  
	
	# 判断source是否是一个图像或视频文件
	is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
	
	# 判断source是否是一个网络链接
	is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
	
	# 判断source是否是一个网络摄像头流或屏幕截图指令或一个有效的URL
	webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
	
	# 判断source是否是一个屏幕截图指令
	screenshot = source.lower().startswith("screen")
	
	# 如果source是一个有效的URL并且指向一个文件，下载这个文件
	if is_url and is_file:
	    source = check_file(source)  
	
	# 创建保存结果的目录，如果存在则覆盖或增量命名
	save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  
	
	# 创建用于保存标签的子目录
	(save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  
	
	# 加载模型并选择设备（CPU或GPU）
	device = select_device(device)
	model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
	
	# 获取模型的步幅、类别名以及模型是否是PyTorch模型
	stride, names, pt = model.stride, model.names, model.pt
	
	# 检查并调整图像尺寸以适应模型的步幅
	imgsz = check_img_size(imgsz, s=stride)  
	
	# 设置batch_size为1，因为通常推理是单张图像进行
	bs = 1  
	
	# 根据source类型选择数据加载方式
	if webcam:
	    # 对于网络摄像头流，检查是否可以显示图像
	    view_img = check_imshow(warn=True)
	    # 加载网络摄像头流数据
	    dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
	    # 确定batch_size
	    bs = len(dataset)
	elif screenshot:
	    # 加载屏幕截图数据
	    dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
	else:
	    # 加载普通图像或视频数据
	    dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
	
	# 初始化视频路径和视频写入器列表
	vid_path, vid_writer = [None] * bs, [None] * bs
	
	# 模型预热
	model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  
	
	# 初始化计数器和时间记录器
	seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
	
	# 遍历数据集中的每一张图片
	for path, im, im0s, vid_cap, s in dataset:
	    # 测量预处理时间
	    with dt[0]:
	        # 将图像转换为Tensor并移至适当设备
	        im = torch.from_numpy(im).to(model.device)
	        # 调整数据类型和归一化
	        im = im.half() if model.fp16 else im.float()  
	        im /= 255  
	        # 扩展维度以匹配batch_size
	        if len(im.shape) == 3:
	            im = im[None]  
	        # 如果模型是XML格式且batch_size大于1，将图像拆分为多个部分
	        if model.xml and im.shape[0] > 1:
	            ims = torch.chunk(im, im.shape[0], 0)
	
	    # 执行推理
	    with dt[1]:
	        # 可视化模式，保存可视化结果
	        visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
	        # 如果模型是XML格式且batch_size大于1，逐个执行推理
	        if model.xml and im.shape[0] > 1:
	            pred = None
	            for image in ims:
	                if pred is None:
	                    pred = model(image, augment=augment, visualize=visualize).unsqueeze(0)
	                else:
	                    pred = torch.cat((pred, model(image, augment=augment, visualize=visualize).unsqueeze(0)), dim=0)
	            pred = [pred, None]
	        # 否则直接执行推理
	        else:
	            pred = model(im, augment=augment, visualize=visualize)
	
	    # 进行非最大值抑制
	    with dt[2]:
	        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
	
	    # 定义CSV文件路径
	    csv_path = save_dir / "predictions.csv"
	
	    # 定义函数将预测数据写入CSV文件
	    def write_to_csv(image_name, prediction, confidence):
	        """将预测数据写入CSV文件，如果文件不存在则创建新文件。"""
	        data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence}
	        with open(csv_path, mode="a", newline="") as f:
	            writer = csv.DictWriter(f, fieldnames=data.keys())
	            if not csv_path.is_file():
	                writer.writeheader()
	            writer.writerow(data)
	
	    # 处理预测结果
	    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)
	
	        # 将路径转换为Path对象
	        p = Path(p)  
	
	        # 构建保存图像的路径
	        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 = Annotator(im0, line_width=line_thickness, example=str(names))
	
	        # 如果有检测结果
	        if len(det):
	            # 对每个类别进行计数
	            for c in det[:, 5].unique():
	                n = (det[:, 5] == c).sum()  
	                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  
	
	            # 写入预测结果到CSV文件
	            if save_csv:
	                write_to_csv(p.name, label, confidence_str)
	
	            # 将检测框坐标从模型输出大小缩放回原图大小
	            det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
	
	            # 将检测结果写入文件或在图像上绘制
	            for *xyxy, conf, cls in reversed(det):
	                c = int(cls)  # 整数类别
	                label = names[c] if hide_conf else f"{names[c]}"  
	                confidence = float(conf)
	                confidence_str = f"{confidence:.2f}"
	
	                # 如果需要保存CSV文件，写入数据
	                if save_csv:
	                    write_to_csv(p.name, label, confidence_str)
	
	                # 如果需要保存文本标签文件，写入数据
	                if save_txt:  
	                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  
	                    line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  
	                    with open(f"{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)
	
	        # 绘制结果
	        im0 = annotator.result()
	
	        # 如果需要显示图像，在窗口中显示
	        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)  
	                cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
	            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 = str(Path(save_path).with_suffix(".mp4"))  
	                    vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
	                vid_writer[i].write(im0)
	
	    # 输出单张图像的推理时间
	    LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
	
	# 输出整体速度统计
	t = tuple(x.t / seen * 1e3 for x in dt)  
	LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}")
	
	# 输出保存结果的信息
	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[0])

1. 归一化操作

python 复制代码 
    # 遍历数据集中的每一张图片
	for path, im, im0s, vid_cap, s in dataset:
	    # 测量预处理时间
	    with dt[0]:
	        # 将图像转换为Tensor并移至适当设备
	        im = torch.from_numpy(im).to(model.device)
	        # 调整数据类型和归一化
	        im = im.half() if model.fp16 else im.float()  
	        im /= 255

这段代码是针对深度学习模型输入预处理的一部分，特别是在使用PyTorch框架时。这里是针对YOLOv5或类似的模型，它说明了如何将输入图像从uint8格式转换为适合模型输入的格式，即fp16（半精度浮点数，16位）或fp32（单精度浮点数，32位）。

代码解析

uint8

uint8类型用于数字图像时，每个像素的颜色通道（如红、绿、蓝）通常使用uint8类型表示，每个通道的值范围从0（黑色）到255（最饱和的颜色）。

精度转换
python 复制代码 
im = im.half() if model.fp16 else im.float()

这段代码检查模型是否支持半精度（fp16）计算。如果model.fp16为True，则im.half()将图像张量从uint8转换为fp16（半精度浮点数）。如果model.fp16为False，则im.float()将图像张量转换为fp32（单精度浮点数）。

注意 ：uint8到fp16或fp32的转换是隐式的，即当从uint8类型转换到浮点类型时，原本的整数值会被转换为相应的浮点数值，但不会改变其数值大小。例如，uint8的255在转换为fp32后仍然是255.0。

归一化
python 复制代码 
im /= 255  # 0 - 255 to 0.0 - 1.0

这行代码将图像张量的像素值从uint8的0到255范围归一化到fp16或fp32的0.0到1.0之间。这是深度学习模型输入预处理中常见的一步，帮助模型在训练和推断时获得更好的数值稳定性，同时使不同亮度和对比度的图像在模型眼中更加"平等"。

这段代码的关键在于它确保了输入图像被适当地格式化和归一化，以供模型进行有效处理。模型是否使用半精度计算取决于模型自身的配置（model.fp16），这通常在模型训练时为了提高计算效率和减少内存使用而设定。归一化步骤则是深度学习图像处理中普遍采用的预处理步骤，确保模型输入的一致性和数值稳定性。

2. 扩展维度

python 复制代码 
# 扩展维度以匹配batch_size
if len(im.shape) == 3:
	im = im[None]

在深度学习中，尤其是使用卷积神经网络（CNN）进行图像处理时，通常需要处理的是一批图像而非单一图像。这是因为现代GPU架构设计为并行处理大量数据，处理一批图像比一次处理一张图像更有效率。因此，即使输入的是单张图像，也需要将其形状转换为适用于模型的批次输入格式。

为什么扩展维度？

当你的图像数据im的形状是三维的（例如，形状为(height, width, channels)），这意味着你只有一个图像。然而，大多数深度学习框架和模型期望输入数据的形状至少是四维的，即(batch_size, height, width, channels)（对于TensorFlow）或(batch_size, channels, height, width)（对于PyTorch）。

在YOLOv5的情况下，模型预期的输入是四维的，即 (batch_size, channels, height, width)。因此，如果你的im是一个单独的图像，它的形状会是 (channels, height, width)，需要在前面增加一个维度来代表batch_size，这样形状就会变成 (1, channels, height, width)。这就是为什么使用im = im[None]来扩展维度，None在这里等价于np.newaxis，它会在数组中插入一个新的轴。

代码解释

python 复制代码 
if len(im.shape) == 3:
    im = im[None]  # 扩展维度，使形状从 (channels, height, width) 变为 (1, channels, height, width)

这条语句检查im的形状，如果它只有三个维度，那么就使用None来扩展其第一个维度，从而匹配模型期望的输入形状。扩展维度是为了将单张图像转换为批次格式，以便模型能够正确处理。这是深度学习实践中一个常见的预处理步骤，尤其在使用CNN进行图像处理时。

三、定义命令行参数

python 复制代码 
def parse_opt():
    parser = argparse.ArgumentParser()  # 创建ArgumentParser对象
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model path(s)')  # 添加权重参数
    parser.add_argument('--source', type=str, default='data/images', help='file/dir/URL/glob, 0 for webcam')  # 添加输入源参数
    parser.add_argument('--data', type=str, default='data/coco128.yaml', help='(optional) dataset.yaml path')  # 添加数据集参数
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')  # 添加图像大小参数
    parser.add_argument('--conf-thres', type=float, default=0.25, help='confidence threshold')  # 添加置信度阈值参数
    parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')  # 添加IoU阈值参数
    parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')  # 添加最大检测数量参数
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')  # 添加设备参数
    parser.add_argument('--view-img', action='store_true', help='show results')  # 添加显示图像参数
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')  # 添加保存文本参数
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')  # 添加保存置信度参数
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')  # 添加保存裁剪框参数
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')  # 添加不保存图像/视频参数
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')  # 添加类别过滤参数
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')  # 添加类别无关的NMS参数
    parser.add_argument('--augment', action='store_true', help='augmented inference')  # 添加增强推理参数
    parser.add_argument('--visualize', action='store_true', help='visualize features')  # 添加可视化特征参数
    parser.add_argument('--update', action='store_true', help='update all models')  # 添加更新模型参数
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')  # 添加项目路径参数
    parser.add_argument('--name', default='exp', help='save results to project/name')  # 添加结果文件夹名称参数
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')  # 添加允许现有项目名称参数
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')  # 添加边界框厚度参数
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')  # 添加隐藏标签参数
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')  # 添加隐藏置信度参数
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')  # 添加半精度推理参数
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')  # 添加OpenCV DNN推理参数
    parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')  # 添加视频帧率步幅参数
    opt = parser.parse_args()  # 解析命令行参数
    return opt  # 返回解析结果

四、主函数

python 复制代码 
def main(opt):
    check_requirements(exclude=('tensorboard', 'thop'))  # 检查运行所需的库
    run(**vars(opt))  # 运行检测程序

if __name__ == "__main__":
    opt = parse_opt()  # 解析命令行参数
    main(opt)  # 运行主程序
相关推荐
奔跑中的小象5 小时前
统信UOS + 天数AI卡部署SGLang服务手册
人工智能·uos·sglang·天数智芯
DevSecOps选型指南5 小时前
中国版Mythos，为何是悬镜安全灵脉CodeAI？
人工智能·安全
Shockang5 小时前
LangGraph 状态机实战
人工智能
刹那芳华19925 小时前
循环神经网络的从零开始实现(RNN)
人工智能·rnn·深度学习
阿童木写作5 小时前
跨境图片翻译工具多合一，批量图片视频字幕翻译加智能抠图
人工智能·python·音视频·语音识别
科技之门5 小时前
AI3D从建模到贴图、绑骨和动画的完整流程怎么做？V2Fun完整工作流指南
人工智能·3d·贴图
宇宙第一小趴菜5 小时前
二、机器学习的应用领域和发展史
人工智能·机器学习
前端开发江鸟6 小时前
我写过 MCP Server，却一直以为 MCP 只有 Tool
人工智能
天国梦6 小时前
自习室智能化升级避坑指南：天学网AI智习室方案实测与选型建议
大数据·人工智能
阿里云云原生6 小时前
可用性从 99.9% 跃升至 99.995%：畅捷通如何用 AI 重塑运维底座？
运维·网络·人工智能
热门推荐
01如何新建文件夹？ 电脑新建文件夹的4种方法02GitHub 镜像站点03国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）04微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新05AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路06Agnes AI 免费 API 接入指南：文本、生图、生视频，一套接口全免费072026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？082026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片0900 Debian字符界面如何支持中文10历年考研数学一、数学二、数学三真题试卷及答案PDF