YoloV5实时推理最短的代码

YoloV5实时推理最简单代码

import cv2
import torch

# 加载YOLOv5模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# 使用CPU或GPU进行推理
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:
        break

    # 进行推理
    results = model(frame)

    # 获取检测结果的图像
    output_img = results.render()[0]

    # 显示图像
    cv2.imshow('YOLOv5', output_img)

    # 检测键盘输入，按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()

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