import cv2
import numpy as np
import os
import time

def read_video_to_memory(video_path):
    """
    读取 MP4 视频并将每一帧存储到内存中
    """
    start_time = time.time()
    cap = cv2.VideoCapture(video_path)
    frames = []
    
    # 获取视频的总帧数和帧率
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    print(total_frames)

    print(fps)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)
    
    cap.release()
    print(time.time() -start_time )
    return frames, fps, total_frames

def generate_thumbnail(frame, max_size=144):
    """
    生成不超过指定大小的缩略图
    """
    height, width = frame.shape[:2]
    
    # 确定缩放比例
    scale = min(max_size / height, max_size / width)
    new_width = int(width * scale)
    new_height = int(height * scale)
    
    # 缩放并返回缩略图
    thumbnail = cv2.resize(frame, (new_width, new_height))
    return thumbnail

def save_video_slice(frames, slice_range, output_file, fps):
    """
    根据指定的帧范围将视频保存为新的片段
    """
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 用于生成 MP4 格式
    height, width = frames[0].shape[:2]
    out = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
    
    start_frame, end_frame = slice_range
    
    # 遍历指定范围的帧并写入到新视频中
    for i in range(start_frame, end_frame):
        out.write(frames[i])
    
    out.release()

def create_video_slice_from_memory(video_path, slice_segments, output_dir):
    """
    从内存中的视频帧创建多个视频片段
    """
    frames, fps, total_frames = read_video_to_memory(video_path)
    
    start_time = time.time()
    # 生成缩略图并显示
    thumbnails = []
    for frame in frames:
        thumbnails.append(generate_thumbnail(frame))

    print(time.time()-start_time)
    
    
    # 根据切片分割视频并保存
    for i, (start_frame, end_frame) in enumerate(slice_segments):
        output_file = os.path.join(output_dir, f"slice_{i+1}.mp4")
        start_time = time.time()-start_time
        save_video_slice(frames, (start_frame, end_frame), output_file, fps)
        print(time.time()-start_time)


# 示例：读取视频并创建多个切片
video_path = "E:/抖音-AI风景/Test/demo.mp4"
slice_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,303)]  # 定义帧区间
output_dir = "output_slices"

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

create_video_slice_from_memory(video_path, slice_segments, output_dir)