yolov8自定义保存

YOLOV8 ultralytics 视频预测代码_yolov8 保存预测视频-CSDN博客
 

import time

import cv2

from ultralytics.models import YOLO
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
# 初始化模型
model = YOLO(r'E:\ultralytics\runs\detect\5种标注扩容train 最好\weights\last.pt') # select your model.pt path
# 打开视频文件
cap = cv2.VideoCapture(0)
# 获取视频帧的维度
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))

# frame_count = 150  # 5秒视频，假设每秒30帧
frame_rate = 30

frames = []
#创建VideoWriter对象
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out1 = cv2.VideoWriter(r'E:\ultralytics-mainmg\new.mp4', fourcc, 20.0, (frame_width, frame_height))
#循环视频帧
save_flag = False  # 保存视频的标志
elapsed_time = 0  # 初始时间为0
target_class_value = 1  # 假设目标类别值为1
start_time = time.time()

while cap.isOpened():
    # 读取某一帧
    success, frame = cap.read()

    if success:

        frames.append(frame)
        # 使用yolov8进行预测
        results = model(frame)
        #可视化结果
        annotated_frame = results[0].plot()
        #将带注释的帧写入视频文件
        out1.write(annotated_frame)
        for r in results:
            boxes = r.boxes  # Boxes object for bbox outputs
            if boxes.cls.numel() != 0:
                # cls_value = boxes.cls.item()
                # print('cls是', cls_value)
                cls_tensor = boxes.cls
                cls_list = cls_tensor.cpu().numpy().tolist()
            else:
                cls_list = [99999]
                print('boxes.cls 中没有元素，无法转换为标量', cls_list)
            end_time = time.time()
            elapsed_time = end_time - start_time
            print('间隔时间', elapsed_time)
            if elapsed_time > 10 and save_flag:  # 重置保存标志
                save_flag = False
            if target_class_value in cls_list and not save_flag:
            #if target_class_value in cls_list:
                timestamp = int(time.time())
                img_name = f'{timestamp}.jpg'
                cv2.imwrite(img_name, frame)
                print(f'保存图片: {img_name}')

                # video_name = f'video_{timestamp}.avi'
                # # 保存前5秒的视频
                # out = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*'DIVX'), frame_rate,
                #                       (frame.shape[1], frame.shape[0]))
                # for i in range(frame_count):
                #     print(i)
                #     out.write(frames[i])
                #
                # out.release()
                # print(f'保存视频: {video_name}')

                # end_time = time.time()
                # elapsed_time = end_time - start_time
                # print('间隔时间',elapsed_time)
               # frame_count = int(elapsed_time * frame_rate)  # 计算已经保存的帧数
                frame_count = len(frames)
                print('帧数',frame_count)

              #  start_frame = max(0, frame_count - frame_rate * 5)  # 判断起始帧
                start_frame = max(0, frame_count - frame_rate * 10)  # 判断起始帧
              #  end_frame = min(frame_count + frame_rate * 5, len(frames))  # 判断结束帧
                end_frame = min(frame_count + frame_rate * 5, len(frames))  # 判断结束帧
                video_name = f'video_{timestamp}.avi'
                out = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*'DIVX'), frame_rate,
                                      (frame.shape[1], frame.shape[0]))
                for i in range(start_frame, end_frame):
                    out.write(frames[i])
                    print(i)
                out.release()
                print(f'保存视频: {video_name}')

                save_flag = True  # 设置保存标志为 True
                frames = []
                start_time = time.time()
            # if elapsed_time > 10 and save_flag:  # 重置保存标志
            #     save_flag = False
    else:
        # 最后结尾中断视频帧循环
        break
#释放读取和写入对象
cap.release()
out.release()