import cv2
from scenedetect import detect, ContentDetector

# 检测视频中的场景
scene_list = detect('demo3.mp4', ContentDetector())

# 存储修改后的场景信息
modified_scene_list = []

# 遍历场景
for i, scene in enumerate(scene_list):
    if i == 0:
        # 第一个场景直接添加
        modified_scene_list.append(scene)
    elif i == len(scene_list) - 1:
        # 最后一个场景不处理
        modified_scene_list.append(scene)
    else:
        # 从第二个场景开始，处理前两帧
        current_scene_start_frame = scene[0].get_frames()
        current_scene_end_frame = scene[1].get_frames()

        # 获取前一个场景
        previous_scene = modified_scene_list[-1]

        # 获取前一个场景的结束帧
        previous_scene_end_frame = previous_scene[1].get_frames()

        # 更新前一个场景的结束帧，将当前场景的前两帧拼接到前一个场景后面
        modified_previous_scene_end_frame = previous_scene_end_frame + 2

        # 更新前一个场景的结束时间
        modified_scene_list[-1] = (previous_scene[0], scene[0])

        # 添加当前场景（去掉前两帧）
        modified_scene_list.append((scene[0], scene[1]))

# 打开原始视频
cap = cv2.VideoCapture('demo3.mp4')
fps = cap.get(cv2.CAP_PROP_FPS)  # 获取帧率
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取帧宽度
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取帧高度

# 对每个场景进行切割并保存为 MP4
for i, scene in enumerate(modified_scene_list):
    start_frame = scene[0].get_frames()
    end_frame = scene[1].get_frames()
    
    # 设置视频读取的起始位置
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    
    # 创建一个 VideoWriter 对象，用于保存当前场景为新的 MP4 文件
    output_filename = f'scene_{i+1}.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'H264')  # 使用 H.264 编码器
    out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_width, frame_height))  # 视频分辨率
    
    # 读取每一帧并写入到输出文件中
    for frame_num in range(start_frame, end_frame):
        ret, frame = cap.read()
        if not ret:
            break
        out.write(frame)

    # 释放 VideoWriter 和 VideoCapture 对象
    out.release()

    print(f'Scene {i+1} has been saved as {output_filename}')

# 释放 VideoCapture 对象
cap.release()