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: break results = 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()
)。