import cv2
import numpy as np
import subprocess

def read_video_frames_to_memory(video_path, frame_rate=30):
    """
    使用 FFmpeg 将视频帧以 PNG 格式直接读取到内存中。
    """
    ffmpeg_command = [
        "ffmpeg", "-i", video_path,
        "-vf", f"fps={frame_rate}",
        "-f", "image2pipe", "-vcodec", "png", "-"
    ]

    frame_size = 1080*1440*3

    # 使用 subprocess 打开 FFmpeg 进程
    process = subprocess.Popen(ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=frame_size**1)

    frames = []
    try:
        while True:
            # 读取单帧数据
            frame_bytes = process.stdout.read(frame_size)  # 每次尝试读取 1MB
            if not frame_bytes:
                break

            # 解码 PNG 数据为 NumPy 数组（使用 OpenCV）
            frame = cv2.imdecode(np.frombuffer(frame_bytes, np.uint8), cv2.IMREAD_UNCHANGED)
            if frame is not None:
                frames.append(frame)
    finally:
        process.stdout.close()
        process.wait()

    return frames

def write_frames_to_mp4(frames, output_path, frame_rate=30):
    """
    从帧数据生成 MP4 文件。
    """
    if not frames:
        print("No frames to write.")
        return

    # 获取帧尺寸 (高度, 宽度)
    frame_height, frame_width = frames[0].shape[:2]
    print(frame_width,frame_height)

    # 定义视频编码器和输出文件
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # MP4 格式编码器
    video_writer = cv2.VideoWriter(output_path, fourcc, frame_rate, (frame_width, frame_height))

    for frame in frames:
        # 检查帧是否是三通道 (RGB/BGR)，如果需要转换格式
        if len(frame.shape) == 2:  # 灰度图
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
        elif frame.shape[2] == 4:  # RGBA
            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)

        video_writer.write(frame)

    video_writer.release()
    print(f"Video saved to {output_path}")

if __name__ == "__main__":
    video_path = "E:/抖音-AI风景/Test/demo.mp4"
    output_path = "output.mp4"

    # 读取视频帧到内存
    frames = read_video_frames_to_memory(video_path)

    # 将帧数据写入 MP4 文件
    write_frames_to_mp4(frames, output_path)