引言
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}') # 显示处理完成的消息和视频保存路径
结语
大家一起努力,一步一个脚印,加油!!!