当前位置:   article > 正文

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

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脚本的解析与简化

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/人工智能uu/article/detail/888874
推荐阅读
相关标签
  

闽ICP备14008679号