Yolo模型可以对视频进行目标检测,但是传统的model.predict方法无法控制播放视频的尺寸。这里我们设计一个Yolo视频检测的封装函数,方便控制视频的尺寸与播放速度。
代码实现:
import cv2
from ultralytics import YOLO
import torch
"""
model_path: 模型文件的路径
video_path: 视频文件的路径
scale: 视频缩放的比例。(0.1-1.0)
title: 视频窗口标题,注意:不支持中文窗口标题
speed: 视频播放的速度(单位毫秒),默认10
"""
def yolo_video_detected(model_path, video_path, scale, title, speed=10):
model_path = model_path # Change this to your YOLOv8 model's path
video_path = video_path # Change this to your video's path
# Load the trained YOLOv8 model
model = YOLO(model_path)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device: %s" % device)
model.to(device)
batch_size = 8
frames_rgb = []
frames = []
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print("Error: Could not open video.")
exit()
# Process video frames
while True:
ret, frame = cap.read()
if not ret:
print("Finished processing video.")
break
width, height, _ = frame.shape
# 视频图片的缩放
new_shape = [int(64 * scale) * int(height / 64), int(64 * scale) * int(width / 64)]
frame = cv2.resize(frame, new_shape)
frames.append(frame)
# YOLOv8 expects RGB images
if len(frames) == batch_size:
with torch.no_grad():
results = model.predict(frames)
# Process each detection
for i, result in enumerate(results):
for box in result.boxes:
print(box.conf)
# 置信度大于0.4
if float(box.conf) > 0.4:
x1, y1, x2, y2 = box.xyxy[0]
# Draw the bounding box on the BGR frame
cv2.rectangle(frames[i], (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
# Add a label above the box
cv2.putText(frames[i], result.names[int(box.cls)], (int(x1), int(y1) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
# 创建窗口,并设置标题为中文
cv2.imshow(title, frames[i])
cv2.waitKey(speed)
if cv2.getWindowProperty(title, cv2.WND_PROP_VISIBLE) < 1:
cap.release()
cv2.destroyAllWindows()
exit()
frames.clear()
frames_rgb.clear()
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
yolo_video_detected('./model/yolov8n.pt', './video/test01.mp4', 0.3, 'xian-power', 10)