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"C:\Users\DELL\Desktop\best.pt"    #best.pt替换成自己的
    video_path = r"C:\Users\DELL\Desktop\a.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)

使用方法:将YOLO训练好的best.pt文件路径和你的视频路径放入代码中,会在可视化结束后输出视频到原视频的同级文件夹

YOLO使用教程:YOLOv8改进专栏|专栏介绍&目录-CSDN博客

具体效果请看:

YOLO本地检测并保存

Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐