一、detect.py
作为YOLOv5模型推理的核心执行文件，detect.py实现了从数据加载到结果输出的完整目标检测流水线。本文只讲代码中最主要的opt内函数的含义，这是detect最核心的东西，至于其他的代码注释我会放在下面，有什么不懂可以评论区提问。
下面对每个命令行参数进行详细介绍，说明它们在检测推理过程中的含义和作用：
--weights
指定模型权重文件的路径（或多个路径），也可以是远程 Triton URL。程序会加载这些权重以构建检测模型。
作用：确定使用哪个预训练模型进行检测。--source
指定输入数据的来源，可以是单个文件、目录、URL、通配符模式、摄像头索引甚至屏幕截图。
作用：决定检测时所使用的输入数据，如图片、视频或实时流。--data
指定数据集的配置文件（通常是 YAML 文件）的路径，该文件包含了类别名称、数据路径等信息。
作用：为模型提供数据集相关的参数，保证类别和数据格式的一致性。--imgsz / --img / --img-size
指定推理时输入图像的尺寸，通常是一个或两个整数（如 640 或 [640, 640]），代表高度和宽度。
作用：设置模型输入尺寸，确保输入大小满足模型的步长要求，并可在一定程度上影响检测速度和精度。--conf-thres
置信度阈值，模型预测的目标如果置信度低于该值将被舍弃，默认是0.5。
作用：过滤低置信度的检测结果，平衡误检与漏检问题。--iou-thres
非极大值抑制（NMS）使用的 IoU（交并比）阈值，决定在有多个重叠框时如何筛选保留，默认是0.45。
作用：降低重叠检测框的冗余，确保最终输出结果更干净。--max-det
每张图像最多保留的检测框数量。
作用：防止在复杂场景中输出过多检测框，控制输出结果的数量。--device
指定推理所使用的设备，可以是 GPU 编号（例如 '0' 或 '0,1,2,3'）或者 'cpu'。
作用：控制模型在何种硬件上运行，影响计算速度和资源利用。--view-img
设定此标志后，检测过程会实时显示带检测框的图像窗口。
作用：方便调试或演示时查看检测效果。--save-txt
如果设置该选项，检测结果将以文本文件形式保存，文本中包含目标类别和位置信息。
作用：便于后续分析、结果统计或与其他工具集成。--save-conf
与--save-txt配合使用，在保存的文本文件中同时记录检测框的置信度。
作用：提供更详细的检测信息，便于评估模型输出。--save-crop
如果启用，将把每个检测框区域裁剪下来并保存为单独的图片。
作用：用于后续目标分析、数据扩充或验证检测结果。--nosave
设置后不保存任何检测结果（无图像或视频输出）。
作用：用于只需要实时检测而无需保存结果时，节省存储资源。--classes
接受一个整数列表，用于指定只检测特定类别。例如：--classes 0或--classes 0 2 3。
作用：过滤检测结果，仅保留用户关心的目标类别。--agnostic-nms
启用类别无关的非极大值抑制，在过滤重叠检测框时忽略类别信息。
作用：在某些场景下可以更好地处理重叠目标，减少类别误判。--augment
启用增强推理，即在推理时使用数据增强技术（如多尺度、翻转等）。
作用：有时能提升检测精度，但可能会牺牲一定的推理速度。--visualize
启用推理过程中特征图或其他中间结果的可视化。
作用：帮助用户了解模型内部处理过程，便于调试和模型改进。--update
启用模型更新操作，通常用于优化或剥离模型中多余部分，解决版本警告等问题。
作用：保持模型文件的最新状态和最佳性能。--project
指定保存检测结果的主目录。
作用：组织和管理检测输出，方便查看和区分不同实验的结果。--name
定义当前检测运行的名称，将作为--project目录下的子文件夹名称。
作用：标识和区分不同的检测实验。--exist-ok
如果目标保存目录已存在，设置此标志后将不会自动递增目录名，而是直接使用现有目录。
作用：控制是否覆盖或递增保存目录，方便连续实验管理。--line-thickness
设置绘制检测框时线条的宽度，单位为像素。
作用：调整检测框在图像上显示的粗细，提升可视化效果。--hide-labels
启用后在输出图像上不显示检测框对应的类别标签。
作用：在某些需求下只需关注检测框的位置而无需文字信息。--hide-conf
启用后在输出图像上不显示检测框的置信度。
作用：使得检测结果更加简洁，适用于只关注类别的场景。--half
启用 FP16 半精度推理，能够加速计算并减少显存占用，但要求硬件支持。
作用：在保证精度的前提下提升推理速度，适用于大规模或实时检测。--dnn
启用 OpenCV 的 DNN 模块进行 ONNX 模型推理，作为一种替代方案。
作用：可能在某些平台上提供更高的推理效率或更好的兼容性。--vid-stride
指定视频处理时的帧采样步长，数值大于1时将跳过若干帧进行处理。
*作用：加速视频检测处理，对于实时检测或长视频非常有用。
总的来说，这些参数允许用户灵活配置检测任务的各个方面，包括模型加载、输入数据类型、推理设置、结果保存和可视化等，以便适应不同的应用场景和硬件条件。
# 获取当前文件的绝对路径
FILE = Path(__file__).resolve()
# 获取当前文件所在的目录（YOLOv5 的根目录）
ROOT = FILE.parents[0] # YOLOv5 root directory
# 如果 ROOT 不在 Python 的搜索路径中，则添加进 sys.path，保证后续可以导入相关模块
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
# 将 ROOT 转换为相对于当前工作目录的相对路径（方便后续处理）
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
# 从 models.common 模块中导入 DetectMultiBackend（支持多后端推理的检测模型封装）
from models.common import DetectMultiBackend
# 导入数据加载相关模块：图片、视频、截图、流媒体等加载方式以及支持的格式
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
# 导入通用工具函数，包括日志、性能计时、文件检测、图像尺寸检查、OpenCV相关操作、路径管理、非极大值抑制等
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_boxes, strip_optimizer, xyxy2xywh)
# 导入绘图工具：用于在图像上绘制检测框、标签、颜色映射等
from utils.plots import Annotator, colors, save_one_box
# 导入 torch 相关工具：设备选择、智能推理模式包装器
from utils.torch_utils import select_device, smart_inference_mode
# 使用装饰器包装 run 函数，启用智能推理模式（比如自动切换评估模式、关闭梯度计算等）
@smart_inference_mode()
def run(
weights=ROOT / 'yolov5s.pt', # 模型权重路径或 Triton URL
source=ROOT / 'data/images', # 输入源，可以是文件、文件夹、URL、摄像头索引等
data=ROOT / 'data/coco128.yaml', # 数据集配置文件路径（包含类别信息、数据集路径等）
imgsz=(640, 640), # 推理时的输入图像尺寸（高, 宽）
conf_thres=0.25, # 置信度阈值，低于该值的检测将被过滤掉
iou_thres=0.45, # 非极大值抑制 (NMS) 的 IoU 阈值
max_det=1000, # 每张图像最多保留的检测数
device='', # 使用的设备，如 'cpu' 或 GPU 序号（例如 '0' 或 '0,1,2,3'）
view_img=False, # 是否实时显示检测结果（用于调试或演示）
save_txt=False, # 是否将检测结果保存为文本文件（保存目标类别和位置信息）
save_conf=False, # 是否在保存的文本中包含置信度
save_crop=False, # 是否保存检测出的目标区域（裁剪图像）
nosave=False, # 如果设置为 True，则不保存检测后的图像或视频
classes=None, # 指定只检测哪些类别，传入类别编号列表（例如 --classes 0 或 --classes 0 2 3）
agnostic_nms=False, # 是否采用类别无关的非极大值抑制
augment=False, # 是否使用数据增强推理（提高鲁棒性）
visualize=False, # 是否可视化模型特征（如特征图等，用于调试）
update=False, # 是否更新模型（比如剥离不必要的模块以加速推理）
project=ROOT / 'runs/detect', # 结果保存的主目录
name='exp', # 本次检测实验的名称（会作为子文件夹名称）
exist_ok=False, # 如果保存目录已存在，是否允许覆盖（不递增目录名）
line_thickness=3, # 绘制检测框的线宽（单位：像素）
hide_labels=False, # 是否隐藏绘制的类别标签
hide_conf=False, # 是否隐藏绘制的置信度数值
half=False, # 是否采用 FP16 半精度推理（加速推理并减少显存占用）
dnn=False, # 是否使用 OpenCV DNN 模块进行 ONNX 推理
vid_stride=1, # 视频帧的采样步长（用于加速处理视频流）
):
# 将 source 转换为字符串（无论输入是 Path 对象还是字符串）
source = str(source)
# 判断是否需要保存检测图像：如果 nosave 为 False 且 source 不是 txt 文件，则保存检测结果
save_img = not nosave and not source.endswith('.txt') # save inference images
# 判断输入是否为单个文件（通过后缀判断是否在支持的图片或视频格式中）
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
# 判断输入是否为 URL（RTSP、RTMP、HTTP 或 HTTPS 开头）
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
# 判断是否为摄像头输入：可以是数字字符串、txt 文件（存放摄像头地址列表）或者 URL（但非单个文件）
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
# 判断是否为屏幕截图输入
screenshot = source.lower().startswith('screen')
# 如果是 URL 且同时被认为是文件，则调用 check_file 下载文件（或验证文件存在性）
if is_url and is_file:
source = check_file(source) # download
# ========== 目录管理 ==========
# 根据 project 目录和实验名称 name 生成唯一的保存目录（如果存在则递增目录名）
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
# 如果需要保存文本结果，则创建 labels 目录，否则创建 save_dir
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# ========== 加载模型 ==========
# 选择运行设备（如 GPU 或 CPU）
device = select_device(device)
# 加载检测模型，支持多后端（如 PyTorch、ONNX、TensorRT 等），并传入相关配置参数
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) # check image size
# ========== 数据加载 ==========
bs = 1 # 初始化批处理大小为1
if webcam:
# 如果是摄像头输入，则检测当前环境是否支持图像显示
view_img = check_imshow(warn=True)
# 使用 LoadStreams 类加载视频流或摄像头数据
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
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
# ========== 模型预热 ==========
# 对模型进行预热，以提高后续推理速度。预热时传入适当的 batch 大小和图像尺寸
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup
# 初始化计数器和计时器，用于记录处理帧数、窗口管理和各阶段时间
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
# 遍历数据集中的每个样本（图片或视频帧）
for path, im, im0s, vid_cap, s in dataset:
# ========= 预处理 =========
with dt[0]:
# 将 numpy 数组转换为 torch 张量，并移动到模型所在的设备上
im = torch.from_numpy(im).to(model.device)
# 如果模型使用半精度，则将数据转换为 fp16，否则使用 fp32（浮点型）
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
# 将像素值归一化到 [0, 1] 区间（原始范围 0 - 255）
im /= 255
# 如果输入只有 3 个维度，则增加 batch 维度
if len(im.shape) == 3:
im = im[None] # expand for batch dim
# ========= 模型推理 =========
with dt[1]:
# 如果需要可视化，则为当前图像创建一个可视化输出目录
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
# 使用加载的模型进行前向推理，得到原始预测结果
pred = model(im, augment=augment, visualize=visualize)
# ========= 非极大值抑制（NMS） =========
with dt[2]:
# 对模型输出进行 NMS 处理，过滤掉低置信度和重叠率高的框
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
# 可选：第二阶段分类器（目前被注释掉）
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# ========= 处理预测结果 =========
for i, det in enumerate(pred): # 针对每个图像（或视频流中的每一帧）
seen += 1 # 统计处理的图像数
if webcam: # 如果输入为摄像头或视频流（批量大于1）
# 分别获取当前帧的路径、图像数据和帧号
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 对象
# 定义保存图像的完整路径（保存检测结果图像，如 im.jpg）
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]] # normalization gain whwh
# 如果需要保存裁剪图像，则复制一份原图用于裁剪，否则直接使用原图
imc = im0.copy() if save_crop else im0
# 初始化 Annotator 对象，用于在图像上绘制检测框和标签，传入绘制线宽和类别示例
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
if len(det):
# 将预测框的坐标从模型输入尺寸映射回原图尺寸
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
# 遍历每个类别，统计该类别的检测个数并添加到日志字符串中
for c in det[:, 5].unique():
n = (det[:, 5] == c).sum() # detections per class
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() # normalized xywh
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:
# 如果在 Linux 平台上，并且该窗口还未创建，则新建窗口以便显示图像
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])
# 显示图像窗口，等待1毫秒（实现实时更新）
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# ========= 保存检测结果 =========
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]
# 强制将保存路径后缀转换为 .mp4
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")
# ========= 处理完所有图像后的总结 =========
# 计算平均每张图像在预处理、推理和 NMS 阶段所耗时间（以毫秒为单位）
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)}' % t)
if save_txt or save_img:
# 如果保存了文本结果，则统计 label 文件数量，并记录保存目录
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 对模型进行更新（剥离冗余部分，避免警告）
strip_optimizer(weights[0])
# ========== 参数解析 ==========
def parse_opt():
# 创建命令行参数解析器
parser = argparse.ArgumentParser()
# 定义各个命令行参数及其默认值和帮助信息：
parser.add_argument('--weights', nargs='+', type=str, default=r'D:\v5\best5.pt', help='model path or triton URL')
parser.add_argument('--source', type=str, default= r'D:\v6\datasets\vein10\val\images', help='file/dir/URL/glob/screen/0(webcam)')
parser.add_argument('--data', type=str, default=ROOT / 'data/vein.yaml', help='(optional) dataset.yaml path')
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
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')
parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
parser.add_argument('--device', default='0', 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: --classes 0, or --classes 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic 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=ROOT / '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')
parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')
# 处理 imgsz 参数，如果只有一个数值，则复制一次使其成为 (h, w)
opt = parser.parse_args()
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
# 打印参数信息（用于调试）
print_args(vars(opt))
return opt
# ========== 主函数 ==========
def main(opt):
# 检查项目所依赖的第三方库是否满足要求（排除 tensorboard 和 thop）
check_requirements(exclude=('tensorboard', 'thop'))
# 将解析得到的参数传入 run 函数开始检测推理
run(**vars(opt))
# 如果该文件作为主程序执行，则解析命令行参数并调用主函数
if __name__ == "__main__":
opt = parse_opt()
main(opt)
扫一扫
1852
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
