Upload extractor.py

extractor.py

@@ -0,0 +1,59 @@
+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)}")