import cv2
import os

def slice_video_by_frames(video_path, frame_segments, output_dir, codec="mp4v", fps=None):
    """
    使用 OpenCV 实现视频按帧切分。

    参数：
    - video_path: 输入视频路径。
    - frame_segments: 帧区间的列表，例如：[(start_frame1, end_frame1), (start_frame2, end_frame2), ...]。
    - output_dir: 输出文件夹路径。
    - codec: 视频编码器，默认是 "mp4v"（适合生成 MP4）。
    - fps: 帧率，默认为输入视频的帧率。
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 打开视频
    video_capture = cv2.VideoCapture(video_path)
    if not video_capture.isOpened():
        raise ValueError(f"无法打开视频文件: {video_path}")

    # 获取视频属性
    original_fps = video_capture.get(cv2.CAP_PROP_FPS)
    frame_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

    if fps is None:
        fps = original_fps

    print(f"视频属性：总帧数={total_frames}, 帧率={original_fps}, 分辨率={frame_width}x{frame_height}")

    for idx, (start_frame, end_frame) in enumerate(frame_segments):
        # 校验帧范围
        if start_frame < 0 or end_frame > total_frames or start_frame >= end_frame:
            print(f"无效的帧区间：({start_frame}, {end_frame})")
            continue

        # 输出文件路径
        output_file = os.path.join(output_dir, f"slice_{idx+1}.mp4")

        # 初始化视频写入器
        fourcc = cv2.VideoWriter_fourcc(*codec)
        video_writer = cv2.VideoWriter(output_file, fourcc, fps, (frame_width, frame_height))

        # 跳转到起始帧
        video_capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

        # 逐帧读取并写入
        for frame_idx in range(start_frame, end_frame):
            ret, frame = video_capture.read()
            if not ret:
                print(f"无法读取帧：{frame_idx}")
                break
            video_writer.write(frame)

        video_writer.release()
        print(f"生成切片: {output_file} ({start_frame} ~ {end_frame})")

    video_capture.release()
    print("视频切分完成。")

# 示例用法
if __name__ == "__main__":
    video_path = "E:/抖音-AI风景/Test/demo.mp4" # 替换为你的输入视频路径
    output_dir = "output"  # 替换为你的输出文件夹路径
    frame_segments = [(0, 60),(60, 79), (79,98),(98,116), (116,135),(135, 153), (153,172), (172,190), (190,209), (209,228), (228,247), (247,266), (266,285),(285,305)]  # 定义帧区间

    slice_video_by_frames(video_path, frame_segments, output_dir)