Spaces:
Running on Zero
Running on Zero
| """Video probing, frame extraction and re-encoding (ffmpeg / PyAV / TurboJPEG).""" | |
| import functools | |
| import re | |
| import shutil | |
| import subprocess | |
| from fractions import Fraction | |
| from pathlib import Path | |
| import av | |
| import cv2 | |
| import imageio.v3 as iio | |
| from turbojpeg import TJPF_RGB, TurboJPEG | |
| def _run_ffmpeg(args): | |
| cmd = ["ffmpeg", "-hide_banner", "-nostdin", "-y", "-loglevel", "error", *map(str, args)] | |
| result = subprocess.run(cmd, capture_output=True, text=True) | |
| if result.returncode != 0: | |
| raise RuntimeError(f"ffmpeg failed: {' '.join(cmd)}\n{result.stderr}") | |
| def extract_frames(video_path, output_dir, quality=2): | |
| """Write every frame to output_dir/%06d.jpg (from 000000.jpg); returns the frame count. | |
| quality: ffmpeg's -q:v for MJPEG, 2 (near-lossless) .. 31; lower is better. | |
| """ | |
| output_dir = Path(output_dir) | |
| if output_dir.exists(): | |
| shutil.rmtree(output_dir) # so that the returned count is exact | |
| output_dir.mkdir(parents=True) | |
| # -fps_mode passthrough: pass frames through as decoded (no duplication/dropping); same as | |
| # the older -vsync 0, which ffmpeg deprecated in 5.1 and newer versions reject | |
| _run_ffmpeg( | |
| ["-i", video_path, "-fps_mode", "passthrough", "-q:v", quality, "-start_number", 0, | |
| output_dir / "%06d.jpg"] | |
| ) # fmt: skip | |
| return sum(1 for p in output_dir.glob("*.jpg") if re.fullmatch(r"\d{6}", p.stem)) | |
| def frame_size(frame_dir): | |
| """(width, height) of the extracted frames.""" | |
| img = cv2.imread(str(Path(frame_dir) / "000000.jpg")) | |
| if img is None: | |
| raise RuntimeError(f"Could not read the first frame in {frame_dir}") | |
| h, w = img.shape[:2] | |
| return w, h | |
| def video_lwh(video_path): | |
| """(num_frames, width, height) from the container (PyAV via imageio).""" | |
| L, H, W, _ = iio.improps(video_path, plugin="pyav").shape | |
| return L, W, H | |
| def video_fps(video_path): | |
| """Average frame rate (frames / duration); robust to VFR and missing rate metadata.""" | |
| with av.open(str(video_path)) as c: | |
| s = c.streams.video[0] | |
| # libav's average_rate is frames / duration | |
| if s.average_rate and s.average_rate > 0: | |
| return float(s.average_rate) | |
| # fallback: demux (no decoding) to count frames and measure the timestamp span; pts can | |
| # be out of order (B-frames), so use min/max | |
| n, pts = 0, [] | |
| for pkt in c.demux(s): | |
| if pkt.size == 0: # flush packet | |
| continue | |
| n += 1 | |
| if pkt.pts is not None: | |
| pts.append(pkt.pts) | |
| if n == 0: | |
| raise ValueError(f"no video frames in {video_path}") | |
| if len(pts) >= 2: | |
| span = (max(pts) - min(pts)) * float(s.time_base) | |
| if span > 0: | |
| return (n - 1) / span | |
| return float(s.base_rate or s.guessed_rate or 0.0) | |
| def _turbojpeg(): | |
| return TurboJPEG() | |
| def read_frame(path): | |
| """JPEG file -> (H, W, 3) RGB uint8 (TurboJPEG; same pixels as cv2.imread + cvtColor).""" | |
| with open(path, "rb") as f: | |
| return _turbojpeg().decode(f.read(), pixel_format=TJPF_RGB) | |
| def nvenc_available(): | |
| """Whether ffmpeg can encode H.264 on the GPU (NVENC); checked with a tiny trial encode, | |
| since a compiled-in encoder can still fail at runtime (no GPU, old driver).""" | |
| listed = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"], capture_output=True, text=True) | |
| if "h264_nvenc" not in listed.stdout: | |
| return False | |
| trial = subprocess.run( | |
| ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-f", "lavfi", | |
| "-i", "color=c=black:s=256x256:d=0.1", "-c:v", "h264_nvenc", "-preset", "p4", | |
| "-f", "null", "-"], | |
| capture_output=True, | |
| ) # fmt: skip | |
| return trial.returncode == 0 | |
| def h264_args(crf): | |
| """ffmpeg output codec args: GPU NVENC if available, else CPU libx264, yuv420p. | |
| NVENC's constant-quality scale runs ~5 points "cheaper" than x264's CRF for the same file | |
| size, hence the offset. | |
| """ | |
| if nvenc_available(): | |
| codec = ["-c:v", "h264_nvenc", "-preset", "p4", "-cq", crf + 5] | |
| else: | |
| codec = ["-c:v", "libx264", "-preset", "veryfast", "-crf", crf] | |
| return [*codec, "-pix_fmt", "yuv420p"] | |
| def encode_frames(frame_dir, out_path, fps, crf=17): | |
| """frame_dir/%06d.jpg -> constant-frame-rate H.264 video with exactly that many frames. | |
| Used to rebuild variable-frame-rate or damaged videos, whose container frame count disagrees | |
| with the frames ffmpeg decodes, so that every reader sees the same frames. | |
| """ | |
| fps = Fraction(fps).limit_denominator(1000) | |
| _run_ffmpeg( | |
| ["-framerate", fps, "-start_number", 0, "-i", Path(frame_dir) / "%06d.jpg", | |
| *h264_args(crf), out_path] | |
| ) # fmt: skip | |
| return out_path | |
| def resample_fps(in_path, out_path, fps=30, crf=17): | |
| """Re-encode a video to a constant frame rate (drops/duplicates frames).""" | |
| _run_ffmpeg(["-i", in_path, "-vf", f"fps={fps}", "-an", *h264_args(crf), out_path]) | |
| return out_path | |