Spaces:
Running on Zero
Running on Zero
File size: 5,610 Bytes
370e047 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Video I/O via FFmpeg: probe, extract frames, trim, reassemble with audio."""
import os
import json
import shutil
import tempfile
import subprocess
from pathlib import Path
from typing import Optional, Generator, Tuple
from dataclasses import dataclass
@dataclass
class VideoInfo:
width: int
height: int
fps: float
total_frames: int
duration: float
has_audio: bool
codec: str
filepath: str
class VideoProcessor:
def __init__(self, config):
self.config = config
self._verify_ffmpeg()
def _verify_ffmpeg(self):
try:
r = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True, timeout=5)
if r.returncode != 0:
raise RuntimeError("FFmpeg not working")
except FileNotFoundError:
raise RuntimeError("FFmpeg not found. Install it (packages.txt: ffmpeg).")
def get_video_info(self, video_path: str) -> VideoInfo:
cmd = ["ffprobe", "-v", "quiet", "-print_format", "json",
"-show_format", "-show_streams", video_path]
r = subprocess.run(cmd, capture_output=True, text=True)
if r.returncode != 0:
raise ValueError(f"Cannot read video: {video_path}")
probe = json.loads(r.stdout)
vstream, has_audio = None, False
for s in probe.get("streams", []):
if s["codec_type"] == "video" and vstream is None:
vstream = s
elif s["codec_type"] == "audio":
has_audio = True
if not vstream:
raise ValueError("No video stream found")
fps_parts = vstream.get("r_frame_rate", "30/1").split("/")
fps = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 and float(fps_parts[1]) else 30.0
nb = int(vstream.get("nb_frames", 0))
dur = float(probe.get("format", {}).get("duration", 0) or 0)
if nb == 0 and dur > 0:
nb = int(dur * fps)
return VideoInfo(int(vstream["width"]), int(vstream["height"]), fps, nb, dur,
has_audio, vstream.get("codec_name", "unknown"), video_path)
def trim(self, video_path: str, start: float, length: float) -> str:
start = max(0.0, float(start)); length = max(0.1, float(length))
out_dir = tempfile.mkdtemp(prefix="trim_")
out_path = os.path.join(out_dir, f"trimmed_{Path(video_path).stem}.mp4")
cmd = ["ffmpeg", "-y", "-ss", str(start), "-i", video_path, "-t", str(length),
"-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
"-c:a", "aac", "-movflags", "+faststart", out_path]
p = subprocess.run(cmd, capture_output=True, text=True)
if p.returncode != 0 or not os.path.exists(out_path):
raise RuntimeError((p.stderr or "trim failed")[-400:])
return out_path
def _tmp(self, name):
d = os.path.join(self.config.video.temp_dir, name)
os.makedirs(d, exist_ok=True)
return d
def extract_frames_to_dir(self, video_path: str, output_dir: str) -> Tuple[str, VideoInfo]:
os.makedirs(output_dir, exist_ok=True)
info = self.get_video_info(video_path)
ext = self.config.video.frame_format
cmd = ["ffmpeg", "-y", "-i", video_path]
if ext == "png":
cmd += ["-compression_level", "3"]
else:
cmd += ["-qscale:v", "2"]
cmd += [os.path.join(output_dir, f"%06d.{ext}")]
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
_, err = p.communicate()
if p.returncode != 0:
raise RuntimeError(f"Frame extraction failed: {err.decode()[-400:]}")
files = sorted(f for f in os.listdir(output_dir) if f.endswith(ext))
info.total_frames = len(files)
return output_dir, info
def assemble_video(self, frames_dir: str, output_path: str, fps: float,
original_video: Optional[str] = None) -> str:
ext = self.config.video.frame_format
temp_video = output_path + ".temp.mp4"
cmd = ["ffmpeg", "-y", "-framerate", str(fps),
"-i", os.path.join(frames_dir, f"%06d.{ext}"),
"-c:v", self.config.video.output_codec,
"-crf", str(self.config.video.output_crf),
"-preset", self.config.video.output_preset,
"-pix_fmt", self.config.video.pixel_format,
"-movflags", "+faststart",
temp_video if original_video else output_path]
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
_, err = p.communicate()
if p.returncode != 0:
raise RuntimeError(f"Encoding failed: {err.decode()[-400:]}")
if original_video:
info = self.get_video_info(original_video)
if info.has_audio:
mux = ["ffmpeg", "-y", "-i", temp_video, "-i", original_video,
"-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0?",
"-shortest", output_path]
pp = subprocess.Popen(mux, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
_, e2 = pp.communicate()
if pp.returncode != 0:
shutil.move(temp_video, output_path)
else:
os.remove(temp_video)
else:
shutil.move(temp_video, output_path)
return output_path
def cleanup_temp(self, temp_dir: str):
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir, ignore_errors=True)
|