Spaces:
Running on L40S
Running on L40S
File size: 2,315 Bytes
9f818c5 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Shared ffprobe-based video metadata helper for the captioning / dataset scripts.
Returns ``fps``, ``duration`` (seconds), ``width``, ``height`` and
``total_frames`` for a video file. Used by both ``caption_from_video.py``
(to fill the structured caption's media fields) and ``captions_to_sft_jsonl.py``
(to build SFT JSONL rows), so the two stay consistent.
"""
import json
import subprocess
from pathlib import Path
_VIDEO_EXTENSIONS = (".mp4", ".mov", ".avi", ".mkv", ".webm")
def probe_video_metadata(video_path: str | Path) -> dict:
"""Return ``{fps, duration, width, height, total_frames}`` via ffprobe.
Raises:
RuntimeError: if ffprobe fails or the file has no video stream.
"""
video_path = str(video_path)
cmd = [
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_streams",
"-show_format",
video_path,
]
result = subprocess.run(cmd, stdin=subprocess.DEVNULL, capture_output=True, text=True)
if result.returncode != 0:
raise RuntimeError(f"ffprobe failed for {video_path}: {result.stderr}")
data = json.loads(result.stdout)
video_stream = next(
(s for s in data.get("streams", []) if s.get("codec_type") == "video"),
None,
)
if video_stream is None:
raise RuntimeError(f"No video stream found in {video_path}")
fps_str = video_stream.get("avg_frame_rate") or video_stream.get("r_frame_rate") or "30/1"
fps_num, fps_den = (fps_str.split("/") + ["1"])[:2]
fps_den_f = float(fps_den) or 1.0
fps = float(fps_num) / fps_den_f
# Duration: prefer the container format duration, fall back to the stream's.
duration = float(data.get("format", {}).get("duration") or video_stream.get("duration") or 0.0)
width = int(video_stream["width"])
height = int(video_stream["height"])
# nb_frames may be absent; fall back to duration * fps.
total_frames = int(video_stream.get("nb_frames") or round(duration * fps))
return {
"fps": fps,
"duration": duration,
"width": width,
"height": height,
"total_frames": total_frames,
}
|