YoloV5实时推理最简单代码
import cv2
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
cap = cv2.VideoCapture(0)while True:ret, frame = cap.read()if not ret:breakresults = model(frame)output_img = results.render()[0]cv2.imshow('YOLOv5', output_img)if cv2.waitKey(1) & 0xFF == ord('q'):break
cap.release()
cv2.destroyAllWindows()
- 导入必要的库:
cv2
:OpenCV库,用于图像和视频处理。torch
:PyTorch库,用于加载和运行YOLOv5模型。
- 加载YOLOv5模型:
- 使用
torch.hub.load
函数从指定的GitHub仓库 'ultralytics/yolov5'
中加载YOLOv5模型 'yolov5s'
(小型版本)。模型将自动下载并加载。
- 确定推理设备:
- 根据系统是否支持CUDA(GPU加速),选择在GPU(‘cuda’)或CPU(‘cpu’)上执行模型推理。
- 打开摄像头:
- 使用
cv2.VideoCapture
打开默认的摄像头(通常是计算机内置摄像头或外部摄像头)。
- 进入主循环:
- 使用
while True
创建一个无限循环,以连续地从摄像头获取帧并进行目标检测。
- 读取摄像头帧:
- 使用
cap.read()
方法读取摄像头的下一帧图像。ret
表示成功与否,frame
包含捕获的图像。
- 进行推理:
- 将读取的帧传递给YOLOv5模型进行目标检测,结果存储在
results
中。
- 获取检测结果的图像:
- 使用
results.render()
方法获取包含检测框和标签的图像,这里使用索引 [0]
表示获取第一帧的检测结果。
- 显示图像:
- 使用
cv2.imshow()
方法在名为’YOLOv5’的窗口中显示检测结果图像。
- 检测键盘输入:
- 使用
cv2.waitKey(1)
检测键盘输入,等待1毫秒。如果按下键盘上的 ‘q’ 键(ASCII码为0x71),则退出循环。
- 释放摄像头并关闭窗口:
- 在循环结束后,释放摄像头资源(
cap.release()
)并关闭显示窗口(cv2.destroyAllWindows()
)。