Ultralytics本身提供了视频的预测代码,但是我并没有找到视频预测结果的存放位置,在查阅了官方网站和CSDN后我发现了它好像只是提供了视频预测过程中进行预测结果显示的代码,没有发现它怎么保存预测后的视频结果,所以寻找了一个对预测后结果进行导出的代码,并对其进行了部分修改。

import cv2
from ultralytics import YOLO
from tqdm import tqdm   # 进度条

# 加载模型
model = YOLO('yolov8n-pose.pt')

# 打开视频文件
video_path = "/root/local/ultralytics-main/MyVideo_1.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/local/ultralytics-main/output.mp4', fourcc, 25.0, (frame_width, frame_height))

# 设置整个视频处理的进度条
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
pbar = tqdm(total=total_frames, desc="Processing video", unit="frames")

# 处理视频帧
for _ in range(total_frames):
    # 读取某一帧
    success, frame = cap.read()
    if success:
        # 使用yolov8进行预测
        results = model(frame)
        # 可视化结果
        annotated_frame = results[0].plot()
        # 将带注释的帧写入视频文件
        out.write(annotated_frame)
        # 更新进度条
        pbar.update(1)
    else:
        # 最后结尾中断视频帧循环
        break

# 若有部分帧未正常打开,进度条是不会达到百分之百的,下面这行代码会让进度条跑满
pbar.update(total_frames - pbar.n)
# 完成视频处理,关闭进度条
pbar.close()

# 释放读取和写入对象
cap.release()
out.release()

参考链接

更多推荐