modelbeast/server/operators/ffmpeg_frames/run.py
2026-07-12 21:05:22 +10:00

60 lines
1.9 KiB
Python

import argparse
import json
import subprocess
import sys
from pathlib import Path
ap = argparse.ArgumentParser()
ap.add_argument("--input", required=True)
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
args = ap.parse_args()
p = json.loads(args.params)
fps = p.get("fps", 2)
dedupe = p.get("dedupe", True)
fmt = p.get("format", "jpg")
max_frames = int(p.get("max_frames", 300))
cull_pct = float(p.get("sharpness_cull_pct", 5))
outdir = Path(args.outdir)
stem = Path(args.input).stem
frames_dir = outdir / f"{stem}_frames"
frames_dir.mkdir(parents=True, exist_ok=True)
vf = f"fps={fps}"
if dedupe:
vf = f"mpdecimate,{vf}"
cmd = ["ffmpeg", "-y", "-i", args.input, "-vf", vf, "-vsync", "vfr"]
if fmt == "jpg":
cmd += ["-q:v", "1"]
cmd += [str(frames_dir / f"%05d.{fmt}")]
print("+", " ".join(cmd), flush=True)
subprocess.run(cmd, check=True, stderr=sys.stdout)
frames = sorted(frames_dir.glob(f"*.{fmt}"))
print(f"extracted {len(frames)} frames", flush=True)
# blur cull: motion-blurred jpgs compress smaller — drop the smallest N%
if fmt == "jpg" and cull_pct > 0 and len(frames) > 20:
by_size = sorted(frames, key=lambda f: f.stat().st_size)
n_cull = int(len(frames) * cull_pct / 100)
for f in by_size[:n_cull]:
f.unlink()
print(f"culled {n_cull} likely-blurry frames", flush=True)
frames = sorted(frames_dir.glob(f"*.{fmt}"))
if len(frames) > max_frames:
keep_every = len(frames) / max_frames
keep = {frames[int(i * keep_every)] for i in range(max_frames)}
for f in frames:
if f not in keep:
f.unlink()
frames = sorted(frames_dir.glob(f"*.{fmt}"))
print(f"thinned to {len(frames)} frames (max_frames={max_frames})", flush=True)
(outdir / "result.json").write_text(json.dumps({
"outputs": [{"path": frames_dir.name,
"meta": {"count": len(frames), "fps": fps, "format": fmt}}]}))
print("done", flush=True)