video_encoder.py

"""
Video Encoder - FFmpeg wrapper for frame-to-WebM encoding
"""

import subprocess
import os
from typing import List
from PIL import Image
import tempfile
import shutil


def encode_frames_to_webm(frames: List[Image.Image], output_path: str, fps: int = 24) -> str:
    """
    Encode a list of PIL Images to a WebM video file.
    
    Args:
        frames: List of PIL Image objects
        output_path: Path for the output WebM file
        fps: Frames per second
    
    Returns:
        Path to the created video file
    """
    if not frames:
        raise ValueError("No frames provided")
    
    # Create temp directory for frames
    temp_dir = tempfile.mkdtemp(prefix="caption_frames_")
    
    try:
        # Save frames as PNG
        width, height = frames[0].size
        for i, frame in enumerate(frames):
            frame_path = os.path.join(temp_dir, f"frame_{i:05d}.png")
            frame.save(frame_path, "PNG")
        
        # FFmpeg command to encode frames to WebM
        # Using VP8 for better compatibility (no alpha needed, green screen)
        cmd = [
            "ffmpeg", "-y",
            "-framerate", str(fps),
            "-i", os.path.join(temp_dir, "frame_%05d.png"),
            "-c:v", "libvpx",  # VP8 encoder
            "-b:v", "2M",
            "-pix_fmt", "yuv420p",
            "-deadline", "realtime",
            "-cpu-used", "8",
            output_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            raise RuntimeError(f"FFmpeg failed: {result.stderr}")
        
        return output_path
        
    finally:
        # Cleanup temp frames
        shutil.rmtree(temp_dir, ignore_errors=True)


def encode_frames_pipe(frames: List[Image.Image], output_path: str, fps: int = 24) -> str:
    """
    Encode frames to WebM using temp files (more reliable than pipe).
    """
    if not frames:
        raise ValueError("No frames provided")
    
    # Use the file-based approach which is more reliable
    return encode_frames_to_webm(frames, output_path, fps)


# Test
if __name__ == "__main__":
    from canvas_renderer import render_frame
    
    # Generate test frames
    words = ["HELLO", "WORLD"]
    frames = []
    
    # Word 0 active for 12 frames, Word 1 active for 12 frames
    for i in range(24):
        active_idx = 0 if i < 12 else 1
        frame = render_frame(words, active_idx, "hormozi")
        frames.append(frame)
    
    # Encode
    encode_frames_pipe(frames, "test_output.webm", fps=24)
    print("Created test_output.webm")