YOLOV8 ultralytics 视频预测代码

ultralytics本身提供了视频的预测代码,但我只发现了它好像只是提供了视频预测过程中进行预测结果显示的代码,没有发现它怎么保存预测后的视频结果,所以写了一个对预测后结果进行导出的代码。

import cv2
from ultralytics import YOLO
# 加载模型
model = YOLO('/root/runs/segment/train4/weights/best.pt')
# 打开视频文件
video_path = "/root/yolov8datasets/VID_20230625_105534.mp4"
cap = cv2.VideoCapture(video_path)
# 获取视频帧的维度
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
#创建VideoWriter对象
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter('/root/yolov8datasets/output2.mp4', fourcc, 20.0, (frame_width, frame_height))
#循环视频帧
while cap.isOpened():
    # 读取某一帧
    success, frame = cap.read()
    if success:
        # 使用yolov8进行预测
        results = model(frame)
        #可视化结果
        annotated_frame = results[0].plot()
        #将带注释的帧写入视频文件
        out.write(annotated_frame)
    else:
        # 最后结尾中断视频帧循环
        break
#释放读取和写入对象
cap.release()
out.release()

你可能感兴趣的:(YOLO,音视频,python)