extractor.py

import os
import shutil
import tempfile
import zipfile
from pathlib import Path
import ffmpeg
import gradio as gr

RESOLUTIONS = {
    "Original": None,
    "512x512": (512, 512),
    "768x768": (768, 768),
    "1024x576": (1024, 576)
}

def extract_frames(video_path, fps=1, resolution=None):
    temp_dir = tempfile.mkdtemp()
    output_dir = Path(temp_dir) / "frames"
    output_dir.mkdir(parents=True, exist_ok=True)

    output_pattern = str(output_dir / "frame_%04d.jpg")

    filters = [f"fps={fps}", "unsharp=5:5:1.0"]
    if resolution:
        filters.append(f"scale={resolution[0]}:{resolution[1]}")
    filter_str = ",".join(filters)

    try:
        (
            ffmpeg
            .input(video_path)
            .output(output_pattern, vf=filter_str, qscale=2)
            .run()
        )
    except ffmpeg.Error as e:
        print("FFmpeg error:", e)
        raise gr.Error("Failed to extract frames with FFmpeg.")

    return output_dir


def zip_frames(frame_dir):
    zip_path = tempfile.NamedTemporaryFile(delete=False, suffix=".zip").name
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for file_path in Path(frame_dir).glob("*.jpg"):
            zipf.write(file_path, arcname=file_path.name)
    return zip_path


def process_video(video_path, fps, resolution_label):
    resolution = RESOLUTIONS[resolution_label]
    try:
        frames_path = extract_frames(video_path, fps, resolution)
        zip_path = zip_frames(frames_path)
        shutil.rmtree(frames_path.parent)
        return zip_path
    except Exception as e:
        print("Error during processing:", str(e))
        raise gr.Error(f"Failed to extract frames: {str(e)}")