【yolov5改编系列之三】yolov5-6.0中detect.py调用摄像头，官方文档内容太多，看着眼花缭乱，我直接给它删掉一大堆。_yolo5魔改

引言

yolov5官方的detect.py文件集成度比较高。一般情况下我们不需要同时实现这么多的功能，比如：检测图像、视频、摄像头等，包括支持pytorch、tensorflow等模型框架。

本节内容是基于yolov5-6.0的detect.py文件改写的。我相信也适用于相近的yolov5版本。

本节代码功能支撑：windows系统、pytorch模型框架、调用摄像头、基于python语言

使用方法：直接用本节代码替换detect.py文件或者在官方detect.py同级目录下新建一个py文件，配置好环境后直接运行。

（记得是要按空格键退出程序哦，不然视频保存就会出现问题）

代码文件.py

import cv2  # OpenCV库，用于图像和视频处理
import numpy as np  # NumPy库，用于数组和矩阵计算
import torch  # PyTorch库，用于深度学习模型
from models.experimental import attempt_load  # 导入YOLOv5模型
from utils.general import check_img_size, non_max_suppression, scale_coords  # 导入辅助函数
from utils.plots import Annotator, colors  # 导入辅助函数
from utils.augmentations import letterbox  # 导入辅助函数
from utils.torch_utils import select_device  # 导入辅助函数

weights = 'yolov5s.pt'  # 模型权重路径
device = '0'  # 设备号，0表示使用第一个GPU；设置为'cpu'，表示使用CPU
save_path = r"runs/detect/webcam.mp4"  # 视频保存路径
img_size = 640  # 图像大小
stride = 32  # 步长
half = False  # 是否使用半精度浮点数减少内存占用，需要GPU支持

device = select_device(device)  # 设置设备
half &= device.type != 'cpu'  # 如果设备为CPU，则禁用半精度模式

# 导入YOLOv5模型
model = attempt_load(weights, map_location=device)
img_size = check_img_size(img_size, s=stride)  # 检查图像大小是否合法
names = model.names  # 类别名称列表

# 打开摄像头
cap = cv2.VideoCapture(0)  # 0表示打开本地摄像头
frame = 0  # 开始处理的帧数

# 获取视频帧率、宽度和高度，设置输出视频的帧率和大小
ret_val, img0 = cap.read()
fps, w, h = 30, img0.shape[1], img0.shape[0]
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

try:
    # 持续处理视频帧直到退出循环
    while True:
        ret_val, img0 = cap.read()  # 读取视频帧
        if not ret_val:
            break  # 如果没有读取到帧，则退出循环

        frame += 1  # 帧数自增
        print(f'video {frame} {save_path}')

        # 对图像进行Padded resize
        img = letterbox(img0, img_size, stride=stride, auto=True)[0]

        # 转换图像格式
        img = img.transpose((2, 0, 1))[::-1]  # HWC转为CHW，BGR转为RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float() / 255.0  # 像素值归一化到[0.0, 1.0]
        img = img[None]  # [h w c] -> [1 h w c]

        # 模型推理
        pred = model(img)[0]  # 获取模型输出
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, max_det=1000)  # 进行非最大抑制

        # 绘制边框和标签
        det = pred[0]  # 检测结果
        annotator = Annotator(img0.copy(), line_width=3, example=str(names))
        if len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()  # 将边框坐标缩放到原始图像大小
            for *xyxy, conf, cls in reversed(det):
                c = int(cls)  # 类别索引
                label = f'{names[c]} {conf:.2f}'  # 类别标签和置信度
                annotator.box_label(xyxy, label, color=colors(c, True))  # 绘制边框和标签

        # 写入视频帧
        im0 = annotator.result()
        vid_writer.write(im0)

        cv2.imshow('image', im0)  # 显示图像
        cv2.waitKey(1)  # 等待按键输入
        if cv2.waitKey(1) == ord(' '):
            # 当按下 ‘空格’ 键时，强制退出循环
            break
finally:
    # 释放视频写入对象和摄像头对象
    vid_writer.release()
    cap.release()

    print(f'Webcam finish, save to {save_path}')  # 显示处理完成的消息和视频保存路径
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

结语
大家一起努力，一步一个脚印，加油！！！

参考资料
链接: YOLOv5的Tricks | 【Trick13】YOLOv5的detect.py脚本的解析与简化