Create app.py

app.py

@@ -0,0 +1,391 @@
+# app.py
+# A Gradio Space that (1) extracts frames with FFmpeg and (2) upscales each frame with Real-ESRGAN.
+# - CPU-friendly (works without GPU). GPU will be faster if available.
+# - Choose extraction mode, then choose upscale model/scale, JPG quality, and get a ZIP of upscaled frames.
+# - Shows both the FFmpeg command and the ESRGAN settings used.
+
+import os
+import io
+import re
+import math
+import json
+import time
+import shutil
+import zipfile
+import tempfile
+import subprocess
+from pathlib import Path
+from typing import List, Optional
+
+import gradio as gr
+
+# ─────────────────────────────────────────────────────────────
+# FFmpeg detection
+# ─────────────────────────────────────────────────────────────
+
+def _which(name: str) -> Optional[str]:
+    from shutil import which
+    return which(name)
+
+FFMPEG = _which("ffmpeg")
+FFPROBE = _which("ffprobe")
+
+if not FFMPEG or not FFPROBE:
+    MISSING_MSG = (
+        "⚠️ FFmpeg not found. Add a 'packages.txt' with just: \nffmpeg\nlibsm6\nlibxext6\nThen restart the Space."
+    )
+else:
+    MISSING_MSG = ""
+
+# ─────────────────────────────────────────────────────────────
+# ffprobe helpers
+# ─────────────────────────────────────────────────────────────
+
+def ffprobe_json(input_path: str) -> dict:
+    if not FFPROBE:
+        return {}
+    cmd = [FFPROBE, "-v", "error", "-print_format", "json", "-show_streams", "-show_format", input_path]
+    res = subprocess.run(cmd, capture_output=True, text=True)
+    if res.returncode != 0:
+        return {}
+    try:
+        return json.loads(res.stdout)
+    except Exception:
+        return {}
+
+
+def parse_video_info(meta: dict) -> dict:
+    info = {"duration": None, "fps": None, "width": None, "height": None}
+    if not meta:
+        return info
+    try:
+        info["duration"] = float(meta.get("format", {}).get("duration", None))
+    except Exception:
+        pass
+    vstreams = [s for s in meta.get("streams", []) if s.get("codec_type") == "video"]
+    if vstreams:
+        v = vstreams[0]
+        rfr = v.get("r_frame_rate") or v.get("avg_frame_rate")
+        if rfr and "/" in rfr:
+            num, den = rfr.split("/")
+            try:
+                num = float(num)
+                den = float(den)
+                if den != 0:
+                    info["fps"] = num / den
+            except Exception:
+                pass
+        info["width"] = v.get("width")
+        info["height"] = v.get("height")
+    return info
+
+# ─────────────────────────────────────────────────────────────
+# FFmpeg command builder
+# ─────────────────────────────────────────────────────────────
+
+def build_ffmpeg_command(
+    input_path: str,
+    mode: str,
+    every_seconds: float,
+    nth_frame: int,
+    exact_fps: float,
+    start_time: str,
+    end_time: str,
+    long_side: int,
+    out_format: str,
+    jpg_quality: int,
+    png_level: int,
+    scene_detect: bool,
+    scene_thresh: float,
+    out_pattern: str,
+) -> List[str]:
+    if not FFMPEG:
+        raise RuntimeError("FFmpeg not available")
+
+    cmd = [FFMPEG, "-y"]
+    if start_time:
+        cmd += ["-ss", start_time]
+    cmd += ["-i", input_path]
+    if end_time:
+        cmd += ["-to", end_time]
+
+    vf = []
+    if mode == "Every N seconds":
+        vf.append(f"fps={max(1e-6, 1.0/float(every_seconds or 1))}")
+    elif mode == "Every Nth frame":
+        vf.append(f"select='not(mod(n,{max(1, int(nth_frame or 1))}))'")
+        vf.append("setpts=N/FRAME_RATE/TB")
+    elif mode == "Exact FPS":
+        vf.append(f"fps={max(1e-6, float(exact_fps or 1))}")
+    elif mode == "All frames":
+        pass
+    else:
+        vf.append("fps=1")
+
+    if scene_detect:
+        vf.append(f"select='gt(scene,{float(scene_thresh)})',showinfo")
+        vf.append("setpts=N/FRAME_RATE/TB")
+
+    if long_side and long_side > 0:
+        vf.append("scale='if(gt(iw,ih),%d,-1)':'if(gt(iw,ih),-1,%d)':force_original_aspect_ratio=decrease" % (long_side, long_side))
+
+    if vf:
+        cmd += ["-vf", ",".join(vf)]
+
+    if out_format == "jpg":
+        cmd += ["-q:v", str(jpg_quality)]
+    elif out_format == "png":
+        cmd += ["-compression_level", str(png_level)]
+
+    cmd += ["-frame_pts", "1", out_pattern]
+    return cmd
+
+# ──────────────────��──────────────────────────────────────────
+# Real-ESRGAN setup (CPU-friendly)
+# ─────────────────────────────────────────────────────────────
+try:
+    from realesrgan import RealESRGANer
+    from basicsr.archs.rrdbnet_arch import RRDBNet
+    _HAVE_REALESRGAN = True
+except Exception:
+    _HAVE_REALESRGAN = False
+
+
+def get_realesrganer(model_name: str, scale: int, tile: int, half: bool, device: str = "cpu"):
+    if not _HAVE_REALESRGAN:
+        raise RuntimeError("realesrgan is not installed. Check requirements.txt")
+    # Model selection
+    if model_name == "x4plus":
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        model_scale = 4
+        netscale = 4
+        model_path = None  # let the lib auto-download
+    elif model_name == "x4plus-anime":
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        model_scale = 4
+        netscale = 4
+        model_path = None
+    elif model_name == "x2plus":
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        model_scale = 2
+        netscale = 2
+        model_path = None
+    else:
+        raise ValueError("Unknown model")
+
+    if scale not in (2, 4):
+        scale = model_scale
+
+    upsampler = RealESRGANer(
+        scale=model_scale,
+        model_path=model_path,
+        model=model,
+        tile=tile,
+        tile_pad=10,
+        pre_pad=0,
+        half=half,
+        device=device,
+    )
+    return upsampler, model_scale
+
+
+def upscale_images(img_paths: List[Path], out_dir: Path, model_name: str, scale: int, tile: int, precision: str, progress=gr.Progress(track_tqdm=True)) -> List[str]:
+    device = "cuda" if os.environ.get("CUDA_VISIBLE_DEVICES") else "cpu"
+    half = (precision == "half") and (device == "cuda")
+
+    upsampler, model_scale = get_realesrganer(model_name, scale, tile, half, device=device)
+
+    out_paths: List[str] = []
+    for i, p in enumerate(img_paths, 1):
+        try:
+            from PIL import Image
+            img = Image.open(p).convert("RGB")
+            output, _ = upsampler.enhance(np.array(img), outscale=scale)
+            out_img = Image.fromarray(output)
+            out_file = out_dir / p.name
+            out_img.save(out_file, quality=95)
+            out_paths.append(str(out_file))
+        except Exception as e:
+            # Skip problematic files, continue
+            print(f"Upscale failed for {p}: {e}")
+        progress(i / max(1, len(img_paths)))
+    return out_paths
+
+# numpy import (used above)
+import numpy as np
+
+# ─────────────────────────────────────────────────────────────
+# Pipeline: extract → upscale → zip
+# ─────────────────────────────────────────────────────────────
+
+def run_pipeline(
+    video: gr.File | None,
+    mode: str,
+    every_seconds: float,
+    nth_frame: int,
+    exact_fps: float,
+    start_time: str,
+    end_time: str,
+    long_side: int,
+    out_format: str,
+    jpg_quality: int,
+    png_level: int,
+    scene_detect: bool,
+    scene_thresh: float,
+    prefix: str,
+    # Upscale
+    do_upscale: bool,
+    model_name: str,
+    scale: int,
+    tile: int,
+    precision: str,
+):
+    if not video or not video.name:
+        return None, None, None, "Upload a video.", ""
+    if not FFMPEG or not FFPROBE:
+        return None, None, None, "FFmpeg missing. See note below.", MISSING_MSG
+    if do_upscale and not _HAVE_REALESRGAN:
+        return None, None, None, "realesrgan is not installed (see requirements.txt)", ""
+
+    # Work dirs
+    work = Path(tempfile.mkdtemp(prefix="vid2up_"))
+    raw_dir = work / "frames_raw"
+    up_dir = work / "frames_upscaled"
+    raw_dir.mkdir(parents=True, exist_ok=True)
+    up_dir.mkdir(parents=True, exist_ok=True)
+
+    # Build and run FFmpeg
+    pattern = str(raw_dir / f"{prefix}_%05d.{out_format}")
+    cmd = build_ffmpeg_command(
+        input_path=video.name,
+        mode=mode,
+        every_seconds=every_seconds,
+        nth_frame=nth_frame,
+        exact_fps=exact_fps,
+        start_time=start_time.strip(),
+        end_time=end_time.strip(),
+        long_side=long_side,
+        out_format=out_format,
+        jpg_quality=jpg_quality,
+        png_level=png_level,
+        scene_detect=scene_detect,
+        scene_thresh=scene_thresh,
+        out_pattern=pattern,
+    )
+
+    cmd_preview = " ".join([s if " " not in s else f'"{s}"' for s in cmd])
+
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+    if proc.returncode != 0:
+        return None, None, None, f"FFmpeg error:\n{proc.stderr}", cmd_preview
+
+    frames = sorted(raw_dir.glob(f"{prefix}_*.{out_format}"))
+    if not frames:
+        return None, None, None, "No frames extracted.", cmd_preview
+
+    # Optionally upscale
+    if do_upscale:
+        up_paths = upscale_images(frames, up_dir, model_name, scale, tile, precision)
+        gallery = up_paths[:60]
+        final_zip_dir = up_dir
+        detail = f"Frames extracted: {len(frames)} | Upscaled: {len(up_paths)} | Model: {model_name} | Scale: x{scale} | Tile: {tile} | Precision: {precision}"
+    else:
+        gallery = [str(p) for p in frames[:60]]
+        final_zip_dir = raw_dir
+        detail = f"Frames extracted: {len(frames)} | (Upscale: off)"
+
+    # Zip
+    zip_path = work / ("upscaled_frames.zip" if do_upscale else "frames.zip")
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for p in sorted(final_zip_dir.glob("*")):
+            zf.write(p, p.name)
+
+    return gallery, str(zip_path), detail, "", cmd_preview
+
+
+# ─────────────────────────────────────────────────────────────
+# UI
+# ─────────────────────────────────────────────────────────────
+
+def build_ui():
+    with gr.Blocks(theme=gr.themes.Soft(), css="""
+    .cf-title { font-size: 1.6rem; font-weight: 800; }
+    .cmdbox textarea { font-family: ui-monospace, Menlo, monospace; font-size: 12px; }
+    """) as demo:
+        gr.Markdown("""
+        <div class="cf-title">Video → JPG → Upscale (FFmpeg + Real-ESRGAN)</div>
+        Upload a video, extract frames, and optionally upscale each image.
+        """)
+
+        with gr.Row():
+            video = gr.File(label="Upload video", file_types=[".mp4", ".mov", ".mkv", ".avi", ".webm", ".m4v"], type="filepath")
+
+        with gr.Accordion("Extraction", open=True):
+            with gr.Row():
+                mode = gr.Dropdown(["Every N seconds", "Every Nth frame", "Exact FPS", "All frames"], value="Every N seconds", label="Mode")
+                every_seconds = gr.Number(value=1.0, label="Every N seconds")
+                nth_frame = gr.Number(value=30, label="Every Nth frame")
+                exact_fps = gr.Number(value=1.0, label="Exact FPS")
+            with gr.Row():
+                start_time = gr.Textbox(value="", label="Start (HH:MM:SS.mmm)")
+                end_time = gr.Textbox(value="", label="End (HH:MM:SS.mmm)")
+                long_side = gr.Number(value=0, label="Resize long side px (0 = none)")
+            with gr.Row():
+                out_format = gr.Dropdown(["jpg", "png"], value="jpg", label="Output format")
+                jpg_quality = gr.Slider(2, 31, value=3, step=1, label="JPG quality (2=best)")
+                png_level = gr.Slider(0, 9, value=2, step=1, label="PNG compression level")
+            with gr.Row():
+                scene_detect = gr.Checkbox(False, label="Scene-change detect")
+                scene_thresh = gr.Slider(0.0, 1.0, value=0.3, step=0.01, label="Scene threshold")
+                prefix = gr.Textbox(value="frame", label="Filename prefix")
+
+        with gr.Accordion("Upscaling", open=True):
+            with gr.Row():
+                do_upscale = gr.Checkbox(True, label="Upscale frames with Real-ESRGAN")
+                model_name = gr.Dropdown(["x4plus", "x4plus-anime", "x2plus"], value="x4plus", label="Model")
+                scale = gr.Dropdown([2, 4], value=4, label="Output scale")
+            with gr.Row():
+                tile = gr.Number(value=0, label="Tile size (0 = auto)»)  # tiling reduces RAM; try 256/512 on GPU, 0 for auto
+                precision = gr.Dropdown(["auto", "half", "full"], value="auto", label="Precision (GPU=half, CPU=full)")
+
+        run_btn = gr.Button("Run: Extract → (Upscale) → ZIP", variant="primary")
+        gallery = gr.Gallery(label="Preview (first 60)", columns=6, height=300)
+        zip_out = gr.File(label="Download ZIP")
+        details = gr.Markdown("Ready.")
+
+        with gr.Accordion("Show FFmpeg command", open=False):
+            cmd_preview = gr.Textbox(label="ffmpeg command", lines=4, elem_classes=["cmdbox"])
+        if MISSING_MSG:
+            gr.Markdown(f"<span style='color:#b45309'>{MISSING_MSG}</span>")
+
+        def _toggle(mode_val, fmt):
+            return (
+                gr.update(visible=(mode_val == "Every N seconds")),
+                gr.update(visible=(mode_val == "Every Nth frame")),
+                gr.update(visible=(mode_val == "Exact FPS")),
+                gr.update(visible=(fmt == "jpg")),
+                gr.update(visible=(fmt == "png")),
+            )
+
+        mode.change(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
+        out_format.change(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
+        demo.load(_toggle, [mode, out_format], [every_seconds, nth_frame, exact_fps, jpg_quality, png_level])
+
+        run_btn.click(
+            run_pipeline,
+            inputs=[
+                video, mode, every_seconds, nth_frame, exact_fps,
+                start_time, end_time, long_side, out_format, jpg_quality, png_level,
+                scene_detect, scene_thresh, prefix,
+                do_upscale, model_name, scale, tile, precision,
+            ],
+            outputs=[gallery, zip_out, details, gr.Textbox(), cmd_preview],
+            api_name="extract_and_upscale",
+        )
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = build_ui()
+    demo.queue().launch()