yolo --- 对视频进行目标检测,实时可视化预测结果,保存预测视频

 

import os

import cv2
from ultralytics import YOLO


def detect_objects_in_video(best_pt_path, video_path, output_video_name):
    output_video_path = video_path.rsplit('.', 1)[0] + '_' + output_video_name + '.mp4'

    model = YOLO(best_pt_path)

    cap = cv2.VideoCapture(video_path)

    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    while cap.isOpened():
        success, frame = cap.read()

        if success:
            results = model(frame)

            annotated_frame = results[0].plot()

            out.write(annotated_frame)

            cv2.imshow('YOLO Detection', annotated_frame)

            if cv2.waitKey(1) & 0xFF == ord('q'):  # 退出循环的话按“q”
                break
        else:
            break

    cap.release()
    out.release()

    cv2.destroyAllWindows()


if __name__ == "__main__":
    best_pt_path = r"D:\yolo11\ultralytics\runs\detect\train3\weights\best.pt"  # best.pt替换成自己的
    video_path = r"C:\Users\Administrator\Desktop\Counter-strike 2 2024.12.02 - 19.58.00.11.mp4"  # 原视频路径
    output_video_name = "out"

    detect_objects_in_video(best_pt_path, video_path, output_video_name)

    output_video_path = video_path.rsplit('.', 1)[0] + '_' + output_video_name + '.mp4'
    os.startfile(output_video_path)

 

posted @ 2024-12-02 23:02  流水灯  阅读(32)  评论(0编辑  收藏  举报