← 返回首页
Yolo目标检测之视频检测封装函数设计
发表时间:2024-12-05 16:02:08
Yolo目标检测之视频检测封装函数设计

Yolo模型可以对视频进行目标检测,但是传统的model.predict方法无法控制播放视频的尺寸。这里我们设计一个Yolo视频检测的封装函数,方便控制视频的尺寸与播放速度。

1.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)