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

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐