#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Video encoding, decoding, and information extraction utilities.
This module provides functionality for working with video files in robot learning
datasets, including frame extraction at specific timestamps, video encoding from
image sequences, and metadata extraction. It supports multiple video backends
for flexible deployment across different platforms.
The module handles the complexity of video codecs, including inter-frame compression
where frames are stored as differences relative to key frames. This requires
loading preceding key frames when accessing specific timestamps, which the module
handles automatically.
Key Features:
- Multiple backends: Supports torchcodec (when available), pyav, and
video_reader backends with automatic fallback.
- Timestamp-based frame extraction: Extracts frames at specific timestamps
with tolerance checking to ensure synchronization.
- Video encoding: Encodes image sequences to video files using ffmpeg with
configurable codecs and quality settings.
- Metadata extraction: Extracts video and audio stream information using
ffprobe.
- HuggingFace integration: Provides VideoFrame feature type for HuggingFace
datasets.
Classes:
VideoFrame
PyArrow-based feature type for HuggingFace datasets containing video
frames with path and timestamp information.
Functions:
Video decoding:
decode_video_frames
Main interface for decoding frames at timestamps with automatic backend selection.
decode_video_frames_torchcodec
Decode frames using torchcodec backend.
decode_video_frames_torchvision
Decode frames using torchvision backends (pyav or video_reader).
Video encoding:
encode_video_frames
Encode a sequence of PNG images into a video file using ffmpeg.
Video information:
get_video_info
Extract video stream metadata (fps, dimensions, codec).
get_audio_info
Extract audio stream metadata (channels, codec, bitrate).
get_video_pixel_channels
Determine pixel channels from pixel format.
get_image_pixel_channels
Determine pixel channels from PIL Image mode.
Backend management:
get_safe_default_codec
Get default codec backend with fallback logic.
Example:
Decode frames at specific timestamps:
>>> frames = decode_video_frames(
... video_path="videos/episode_0.mp4",
... timestamps=[0.1, 0.2, 0.3],
... tolerance_s=1e-4,
... backend="torchcodec"
... )
Encode images to video:
>>> encode_video_frames(
... imgs_dir="images/episode_0",
... video_path="videos/episode_0.mp4",
... fps=30,
... vcodec="libsvtav1"
... )
Get video information:
>>> info = get_video_info("videos/episode_0.mp4")
>>> print(f"FPS: {info['video.fps']}, Resolution: {info['video.width']}x{info['video.height']}")
"""
import functools
import importlib
import json
import logging
import math
import shutil
import subprocess
import warnings
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, ClassVar
import pyarrow as pa
import torch
import torchvision
from datasets.features.features import register_feature
from PIL import Image
@functools.cache
def _load_torchcodec_videodecoder() -> tuple[type | None, BaseException | None]:
"""Probe whether ``torchcodec.decoders.VideoDecoder`` is usable on this host.
Returns a ``(VideoDecoder, None)`` pair on success, or ``(None, exc)`` if
either the Python package is missing or its compiled extension fails to
load (e.g. no compatible system FFmpeg). Carrying the original exception
forward lets direct callers of the torchcodec decode path surface the root
cause in their own raise, even when the probe's warning has been dropped
(e.g. probe fires before ``init_logging``) or rotated out of a long-lived
log. Wide except: ``find_spec`` can itself raise ``ValueError`` on broken
installs, and the extension import surfaces failures as either
``ImportError`` or ``RuntimeError``. Cached via ``functools.cache`` so the
probe runs once per process and is thread-safe; tests can re-probe with
``_load_torchcodec_videodecoder.cache_clear()``.
"""
try:
if importlib.util.find_spec("torchcodec") is None:
exc = ModuleNotFoundError("torchcodec is not installed on this platform.")
logging.warning("%s Falling back to 'pyav' as the default decoder.", exc)
return None, exc
from torchcodec.decoders import VideoDecoder
except (ImportError, RuntimeError, ValueError) as e:
logging.warning(
"'torchcodec' failed to load (%s); falling back to 'pyav' as the default decoder.",
e,
)
return None, e
return VideoDecoder, None
@functools.cache
def _warn_explicit_torchcodec_unloadable() -> None:
"""Emit the explicit-backend-fallback warning at most once per process.
Without this guard the warning fires inside ``decode_video_frames``, which
sits in ``LeRobotDataset._query_videos`` and is invoked O(num_video_keys)
times per dataloader step. A real run logs tens of thousands of copies.
"""
logging.warning("Caller requested backend='torchcodec' but it failed to load; using 'pyav' instead.")
[docs]
def get_safe_default_codec() -> str:
"""Get the default video codec backend.
Returns ``"torchcodec"`` when the compiled extension actually loads on this
host (requires a compatible system FFmpeg), else ``"pyav"``. The underlying
probe is cached in ``_load_torchcodec_videodecoder``; this wrapper is left
uncached so that ``dataclasses.field(default_factory=...)`` callers get the
fresh-instance semantics they expect from ``default_factory``.
"""
return "torchcodec" if _load_torchcodec_videodecoder()[0] is not None else "pyav"
_safe_encoding_vcodec: str | None = None
[docs]
def get_safe_encoding_vcodec() -> str:
"""Return an ffmpeg encoding codec that is available on this system.
Prefers libsvtav1; falls back to libx264 if libsvtav1 is not available
(e.g. ffmpeg built without SVT-AV1).
Returns:
"libsvtav1" or "libx264".
"""
global _safe_encoding_vcodec
if _safe_encoding_vcodec is not None:
return _safe_encoding_vcodec
ffmpeg_path = shutil.which("ffmpeg")
if ffmpeg_path is None:
logging.warning("ffmpeg not found in PATH, using 'libx264' for video encoding")
_safe_encoding_vcodec = "libx264"
return _safe_encoding_vcodec
try:
result = subprocess.run(
[ffmpeg_path, "-encoders"],
capture_output=True,
text=True,
timeout=5,
)
if result.returncode == 0 and "libsvtav1" in (result.stdout or ""):
_safe_encoding_vcodec = "libsvtav1"
else:
_safe_encoding_vcodec = "libx264"
logging.warning("ffmpeg encoder 'libsvtav1' is not available, using 'libx264' for video encoding")
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Could not detect ffmpeg encoders (%s), using 'libx264'", e)
_safe_encoding_vcodec = "libx264"
return _safe_encoding_vcodec
[docs]
def decode_video_frames(
video_path: Path | str,
timestamps: list[float],
tolerance_s: float,
backend: str | None = None,
) -> torch.Tensor:
"""
Decodes video frames using the specified backend.
Args:
video_path (Path): Path to the video file.
timestamps (list[float]): List of timestamps to extract frames.
tolerance_s (float): Allowed deviation in seconds for frame retrieval.
backend (str, optional): Backend to use for decoding. Defaults to "torchcodec" when available in the platform; otherwise, defaults to "pyav".
Returns:
torch.Tensor: Decoded frames.
Currently supports torchcodec on cpu and pyav. If the caller explicitly
requests ``backend="torchcodec"`` but the host's torchcodec extension
cannot load (e.g. missing system FFmpeg, ABI mismatch), this function
silently downgrades to ``"pyav"`` with a once-per-process warning. The two
decoders pick frames using non-equivalent logic (pts seek + argmin vs.
``round(ts * average_fps)`` indexing), so the decoded pixels for the same
``(video_path, timestamps)`` pair can differ within ``tolerance_s``;
callers needing strict backend selection should check
``get_safe_default_codec() == "torchcodec"`` before dispatching.
"""
if backend is None:
backend = get_safe_default_codec()
if backend == "torchcodec":
if _load_torchcodec_videodecoder()[0] is None:
_warn_explicit_torchcodec_unloadable()
backend = "pyav"
else:
return decode_video_frames_torchcodec(video_path, timestamps, tolerance_s)
if backend in ("pyav", "video_reader"):
return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
raise ValueError(f"Unsupported video backend: {backend}")
[docs]
def decode_video_frames_torchvision(
video_path: Path | str,
timestamps: list[float],
tolerance_s: float,
backend: str = "pyav",
log_loaded_timestamps: bool = False,
) -> torch.Tensor:
"""Loads frames associated to the requested timestamps of a video
The backend can be either "pyav" (default) or "video_reader".
"video_reader" requires installing torchvision from source, see:
https://github.com/pytorch/vision/blob/main/torchvision/csrc/io/decoder/gpu/README.rst
(note that you need to compile against ffmpeg<4.3)
While both use cpu, "video_reader" is supposedly faster than "pyav" but requires additional setup.
For more info on video decoding, see `benchmark/video/README.md`
See torchvision doc for more info on these two backends:
https://pytorch.org/vision/0.18/index.html?highlight=backend#torchvision.set_video_backend
Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
the encoder stores a reference frame (or a key frame) and subsequent frames as differences relative to
that key frame. As a consequence, to access a requested frame, we need to load the preceding key frame,
and all subsequent frames until reaching the requested frame. The number of key frames in a video
can be adjusted during encoding to take into account decoding time and video size in bytes.
"""
video_path = str(video_path)
# set backend
keyframes_only = False
torchvision.set_video_backend(backend)
if backend == "pyav":
keyframes_only = True # pyav doesnt support accuracte seek
# set a video stream reader
# TODO(rcadene): also load audio stream at the same time
reader = torchvision.io.VideoReader(video_path, "video")
# set the first and last requested timestamps
# Note: previous timestamps are usually loaded, since we need to access the previous key frame
first_ts = min(timestamps)
last_ts = max(timestamps)
# access closest key frame of the first requested frame
# Note: closest key frame timestamp is usually smaller than `first_ts` (e.g. key frame can be the first frame of the video)
# for details on what `seek` is doing see: https://pyav.basswood-io.com/docs/stable/api/container.html?highlight=inputcontainer#av.container.InputContainer.seek
reader.seek(first_ts, keyframes_only=keyframes_only)
# load all frames until last requested frame
loaded_frames = []
loaded_ts = []
for frame in reader:
current_ts = frame["pts"]
if log_loaded_timestamps:
logging.info(f"frame loaded at timestamp={current_ts:.4f}")
loaded_frames.append(frame["data"])
loaded_ts.append(current_ts)
if current_ts >= last_ts:
break
if backend == "pyav":
reader.container.close()
reader = None
query_ts = torch.tensor(timestamps)
loaded_ts = torch.tensor(loaded_ts)
# compute distances between each query timestamp and timestamps of all loaded frames
dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
min_, argmin_ = dist.min(1)
is_within_tol = min_ < tolerance_s
assert is_within_tol.all(), (
f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
"It means that the closest frame that can be loaded from the video is too far away in time."
"This might be due to synchronization issues with timestamps during data collection."
"To be safe, we advise to ignore this item during training."
f"\nqueried timestamps: {query_ts}"
f"\nloaded timestamps: {loaded_ts}"
f"\nvideo: {video_path}"
f"\nbackend: {backend}"
)
# get closest frames to the query timestamps
closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
closest_ts = loaded_ts[argmin_]
if log_loaded_timestamps:
logging.info(f"{closest_ts=}")
# convert to the pytorch format which is float32 in [0,1] range (and channel first)
closest_frames = closest_frames.type(torch.float32) / 255
assert len(timestamps) == len(closest_frames)
return closest_frames
[docs]
def decode_video_frames_torchcodec(
video_path: Path | str,
timestamps: list[float],
tolerance_s: float,
device: str = "cpu",
log_loaded_timestamps: bool = False,
) -> torch.Tensor:
"""Loads frames associated with the requested timestamps of a video using torchcodec.
Note: Setting device="cuda" outside the main process, e.g. in data loader workers, will lead to CUDA initialization errors.
Note: Video benefits from inter-frame compression. Instead of storing every frame individually,
the encoder stores a reference frame (or a key frame) and subsequent frames as differences relative to
that key frame. As a consequence, to access a requested frame, we need to load the preceding key frame,
and all subsequent frames until reaching the requested frame. The number of key frames in a video
can be adjusted during encoding to take into account decoding time and video size in bytes.
"""
VideoDecoder, probe_exc = _load_torchcodec_videodecoder() # noqa: N806
if VideoDecoder is None:
raise ImportError(
f"torchcodec is required but not loadable on this host: {probe_exc!r}"
) from probe_exc
# initialize video decoder
decoder = VideoDecoder(video_path, device=device, seek_mode="exact")
loaded_frames = []
loaded_ts = []
# get metadata for frame information
metadata = decoder.metadata
average_fps = metadata.average_fps
num_frames = metadata.num_frames
# convert timestamps to frame indices
frame_indices = [round(ts * average_fps) for ts in timestamps]
# Clamp to the decodable range [0, num_frames - 1]. A query timestamp can map
# past the final frame when a consolidated video was encoded with marginally
# fewer frames than its metadata implies, or when ``average_fps`` rounds up
# near the end of a long video. ``torchcodec``'s ``get_frames_at`` raises a
# hard ``IndexError`` on an out-of-range index, whereas the pyav path returns
# the nearest decoded frame; clamping reproduces that nearest-frame behavior.
# The tolerance check below still rejects a frame that is genuinely too far
# from its query timestamp (and ``retry_random_on_failure`` then skips that
# item), so this only rescues the off-by-rounding boundary case.
if num_frames:
frame_indices = [min(max(int(fi), 0), num_frames - 1) for fi in frame_indices]
# retrieve frames based on indices
frames_batch = decoder.get_frames_at(indices=frame_indices)
for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=False):
loaded_frames.append(frame)
loaded_ts.append(pts.item())
if log_loaded_timestamps:
logging.info(f"Frame loaded at timestamp={pts:.4f}")
query_ts = torch.tensor(timestamps)
loaded_ts = torch.tensor(loaded_ts)
# compute distances between each query timestamp and loaded timestamps
dist = torch.cdist(query_ts[:, None], loaded_ts[:, None], p=1)
min_, argmin_ = dist.min(1)
is_within_tol = min_ < tolerance_s
assert is_within_tol.all(), (
f"One or several query timestamps unexpectedly violate the tolerance ({min_[~is_within_tol]} > {tolerance_s=})."
"It means that the closest frame that can be loaded from the video is too far away in time."
"This might be due to synchronization issues with timestamps during data collection."
"To be safe, we advise to ignore this item during training."
f"\nqueried timestamps: {query_ts}"
f"\nloaded timestamps: {loaded_ts}"
f"\nvideo: {video_path}"
)
# get closest frames to the query timestamps
closest_frames = torch.stack([loaded_frames[idx] for idx in argmin_])
closest_ts = loaded_ts[argmin_]
if log_loaded_timestamps:
logging.info(f"{closest_ts=}")
# convert to float32 in [0,1] range (channel first)
closest_frames = closest_frames.type(torch.float32) / 255
assert len(timestamps) == len(closest_frames)
return closest_frames
[docs]
def encode_video_frames(
imgs_dir: Path | str,
video_path: Path | str,
fps: int,
vcodec: str = "libsvtav1",
pix_fmt: str = "yuv420p",
g: int | None = 2,
crf: int | None = 30,
fast_decode: int = 0,
log_level: str | None = "error",
overwrite: bool = False,
) -> None:
"""Encode a sequence of images into a video file using ffmpeg.
Args:
imgs_dir: Directory containing sequentially numbered PNG frames
(frame_000000.png, frame_000001.png, etc.).
video_path: Output path for the encoded video file.
fps: Frames per second for the output video.
vcodec: Video codec to use. Defaults to "libsvtav1".
pix_fmt: Pixel format. Defaults to "yuv420p".
g: GOP (Group of Pictures) size. Defaults to 2.
crf: Constant Rate Factor for quality control. Defaults to 30.
fast_decode: Fast decode parameter for libsvtav1. Defaults to 0.
log_level: FFmpeg log level. Defaults to "error".
overwrite: Whether to overwrite existing video file. Defaults to False.
Raises:
OSError: If video encoding fails or output file is not created.
Note:
More info on ffmpeg arguments tuning on `benchmark/video/README.md`
"""
video_path = Path(video_path)
imgs_dir = Path(imgs_dir)
video_path.parent.mkdir(parents=True, exist_ok=True)
ffmpeg_path = shutil.which("ffmpeg")
if ffmpeg_path is None:
raise OSError("ffmpeg not found in PATH. Install ffmpeg to encode videos.")
if vcodec == "libsvtav1":
vcodec = get_safe_encoding_vcodec()
ffmpeg_args = OrderedDict(
[
("-f", "image2"),
("-r", str(fps)),
("-i", str(imgs_dir / "frame_%06d.png")),
("-vcodec", vcodec),
("-pix_fmt", pix_fmt),
]
)
if g is not None:
ffmpeg_args["-g"] = str(g)
if crf is not None:
ffmpeg_args["-crf"] = str(crf)
if fast_decode:
key = "-svtav1-params" if vcodec == "libsvtav1" else "-tune"
value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
ffmpeg_args[key] = value
if log_level is not None:
ffmpeg_args["-loglevel"] = str(log_level)
ffmpeg_args = [item for pair in ffmpeg_args.items() for item in pair]
if overwrite:
ffmpeg_args.append("-y")
ffmpeg_cmd = [ffmpeg_path] + ffmpeg_args + [str(video_path)]
# redirect stdin to subprocess.DEVNULL to prevent reading random keyboard inputs from terminal
subprocess.run(ffmpeg_cmd, check=True, stdin=subprocess.DEVNULL)
if not video_path.exists():
raise OSError(
f"Video encoding did not work. File not found: {video_path}. "
f"Try running the command manually to debug: `{''.join(ffmpeg_cmd)}`"
)
[docs]
@dataclass
class VideoFrame:
# TODO(rcadene, lhoestq): move to Hugging Face `datasets` repo
"""
Provides a type for a dataset containing video frames.
Example:
```python
data_dict = [{"image": {"path": "videos/episode_0.mp4", "timestamp": 0.3}}]
features = {"image": VideoFrame()}
Dataset.from_dict(data_dict, features=Features(features))
```
"""
pa_type: ClassVar[Any] = pa.struct({"path": pa.string(), "timestamp": pa.float32()})
_type: str = field(default="VideoFrame", init=False, repr=False)
def __call__(self):
return self.pa_type
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
"'register_feature' is experimental and might be subject to breaking changes in the future.",
category=UserWarning,
)
# to make VideoFrame available in HuggingFace `datasets`
register_feature(VideoFrame, "VideoFrame")
[docs]
def get_audio_info(video_path: Path | str) -> dict:
"""Extract audio stream information from a video file using ffprobe.
Args:
video_path: Path to the video file.
Returns:
Dictionary containing audio information:
- has_audio: Boolean indicating if audio stream exists.
- audio.channels: Number of audio channels (if available).
- audio.codec: Audio codec name (if available).
- audio.bit_rate: Bit rate in bits per second (if available).
- audio.sample_rate: Sample rate in Hz (if available).
- audio.bit_depth: Bit depth (if available).
- audio.channel_layout: Channel layout (if available).
Raises:
RuntimeError: If ffprobe command fails.
"""
ffprobe_audio_cmd = [
"ffprobe",
"-v",
"error",
"-select_streams",
"a:0",
"-show_entries",
"stream=channels,codec_name,bit_rate,sample_rate,bit_depth,channel_layout,duration",
"-of",
"json",
str(video_path),
]
result = subprocess.run(ffprobe_audio_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
if result.returncode != 0:
raise RuntimeError(f"Error running ffprobe: {result.stderr}")
info = json.loads(result.stdout)
audio_stream_info = info["streams"][0] if info.get("streams") else None
if audio_stream_info is None:
return {"has_audio": False}
# Return the information, defaulting to None if no audio stream is present
return {
"has_audio": True,
"audio.channels": audio_stream_info.get("channels", None),
"audio.codec": audio_stream_info.get("codec_name", None),
"audio.bit_rate": int(audio_stream_info["bit_rate"]) if audio_stream_info.get("bit_rate") else None,
"audio.sample_rate": int(audio_stream_info["sample_rate"])
if audio_stream_info.get("sample_rate")
else None,
"audio.bit_depth": audio_stream_info.get("bit_depth", None),
"audio.channel_layout": audio_stream_info.get("channel_layout", None),
}
[docs]
def get_video_info(video_path: Path | str) -> dict:
"""Extract video stream information from a video file using ffprobe.
Args:
video_path: Path to the video file.
Returns:
Dictionary containing video and audio information:
- video.fps: Frames per second.
- video.height: Video height in pixels.
- video.width: Video width in pixels.
- video.channels: Number of pixel channels.
- video.codec: Video codec name.
- video.pix_fmt: Pixel format.
- video.is_depth_map: Whether video is a depth map.
- Plus all fields from get_audio_info().
Raises:
RuntimeError: If ffprobe command fails.
"""
ffprobe_video_cmd = [
"ffprobe",
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=r_frame_rate,width,height,codec_name,nb_frames,duration,pix_fmt",
"-of",
"json",
str(video_path),
]
result = subprocess.run(ffprobe_video_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
if result.returncode != 0:
raise RuntimeError(f"Error running ffprobe: {result.stderr}")
info = json.loads(result.stdout)
video_stream_info = info["streams"][0]
# Calculate fps from r_frame_rate
r_frame_rate = video_stream_info["r_frame_rate"]
num, denom = map(int, r_frame_rate.split("/"))
fps = num / denom
pixel_channels = get_video_pixel_channels(video_stream_info["pix_fmt"])
video_info = {
"video.fps": fps,
"video.height": video_stream_info["height"],
"video.width": video_stream_info["width"],
"video.channels": pixel_channels,
"video.codec": video_stream_info["codec_name"],
"video.pix_fmt": video_stream_info["pix_fmt"],
"video.is_depth_map": False,
**get_audio_info(video_path),
}
return video_info
[docs]
def get_video_pixel_channels(pix_fmt: str) -> int:
"""Determine the number of pixel channels from a pixel format string.
Args:
pix_fmt: Pixel format string (e.g., "yuv420p", "rgb24").
Returns:
Number of channels (1, 3, or 4).
Raises:
ValueError: If pixel format is unknown.
"""
if "gray" in pix_fmt or "depth" in pix_fmt or "monochrome" in pix_fmt:
return 1
elif "rgba" in pix_fmt or "yuva" in pix_fmt:
return 4
elif "rgb" in pix_fmt or "yuv" in pix_fmt:
return 3
else:
raise ValueError("Unknown format")
[docs]
def get_image_pixel_channels(image: Image) -> int:
"""Determine the number of pixel channels from a PIL Image mode.
Args:
image: PIL Image object.
Returns:
Number of channels (1, 2, 3, or 4).
Raises:
ValueError: If image mode is unknown.
"""
if image.mode == "L":
return 1 # Grayscale
elif image.mode == "LA":
return 2 # Grayscale + Alpha
elif image.mode == "RGB":
return 3 # RGB
elif image.mode == "RGBA":
return 4 # RGBA
else:
raise ValueError("Unknown format")
[docs]
def resample_and_trim_video(
input_path: Path | str,
output_path: Path | str,
target_fps: int,
num_frames: int,
vcodec: str = "libsvtav1",
pix_fmt: str = "yuv420p",
g: int | None = 2,
crf: int | None = 30,
fast_decode: int = 0,
log_level: str | None = "error",
overwrite: bool = False,
start_time: float | None = None,
) -> None:
"""Resample a video to the target FPS and trim it to exactly *num_frames* frames.
This is used to attach a pre-recorded MP4 to a :class:`LeRobotDataset`
episode whose non-visual observations have already been saved. The source
video is re-encoded at ``target_fps`` and cropped so that the output
contains exactly ``num_frames`` frames (i.e. a duration of
``num_frames / target_fps`` seconds, starting from ``start_time`` seconds
into the input, or from the beginning if ``start_time`` is ``None``).
Args:
input_path: Path to the source video file.
output_path: Path where the resampled/trimmed video will be written.
target_fps: Desired output frames per second.
num_frames: Exact number of frames the output video must contain.
vcodec: Video codec to use. Defaults to "libsvtav1".
pix_fmt: Pixel format. Defaults to "yuv420p".
g: GOP (Group of Pictures) size. Defaults to 2.
crf: Constant Rate Factor for quality control. Defaults to 30.
fast_decode: Fast decode parameter for libsvtav1. Defaults to 0.
log_level: FFmpeg log level. Defaults to "error".
overwrite: Whether to overwrite an existing output file. Defaults to False.
start_time: Optional start offset in seconds into the source video.
When provided, ``-ss`` is placed before ``-i`` for fast input
seeking. Defaults to None (start from the beginning).
Raises:
FileNotFoundError: If ``input_path`` does not exist.
OSError: If ffmpeg is not found or the output file is not produced.
ValueError: If the resulting video does not have the expected number of
frames (tolerance of ±1 frame to account for codec rounding).
"""
input_path = Path(input_path)
output_path = Path(output_path)
if num_frames <= 0:
raise ValueError(f"`num_frames` must be > 0, got {num_frames}.")
if target_fps <= 0:
raise ValueError(f"`target_fps` must be > 0, got {target_fps}.")
if start_time is not None and (not math.isfinite(start_time) or start_time < 0):
raise ValueError(f"`start_time` must be a finite non-negative number, got {start_time}.")
if not input_path.is_file():
raise FileNotFoundError(f"Input video not found: {input_path}")
output_path.parent.mkdir(parents=True, exist_ok=True)
ffmpeg_path = shutil.which("ffmpeg")
if ffmpeg_path is None:
raise OSError("ffmpeg not found in PATH. Install ffmpeg to process videos.")
if vcodec == "libsvtav1":
vcodec = get_safe_encoding_vcodec()
# Duration to keep: num_frames / target_fps
duration = num_frames / target_fps
input_args: list[tuple[str, str]] = []
if start_time is not None:
input_args.append(("-ss", f"{start_time:.9f}"))
input_args.append(("-i", str(input_path)))
ffmpeg_args = OrderedDict(
input_args
+ [
("-t", f"{duration:.6f}"),
("-r", str(target_fps)),
("-vcodec", vcodec),
("-pix_fmt", pix_fmt),
]
)
if g is not None:
ffmpeg_args["-g"] = str(g)
if crf is not None:
ffmpeg_args["-crf"] = str(crf)
if fast_decode:
key = "-svtav1-params" if vcodec == "libsvtav1" else "-tune"
value = f"fast-decode={fast_decode}" if vcodec == "libsvtav1" else "fastdecode"
ffmpeg_args[key] = value
if log_level is not None:
ffmpeg_args["-loglevel"] = str(log_level)
ffmpeg_args_list = [item for pair in ffmpeg_args.items() for item in pair]
if overwrite:
ffmpeg_args_list.append("-y")
ffmpeg_cmd = [ffmpeg_path] + ffmpeg_args_list + [str(output_path)]
subprocess.run(ffmpeg_cmd, check=True, stdin=subprocess.DEVNULL)
if not output_path.exists():
raise OSError(
f"Video resampling/trimming did not work. File not found: {output_path}. "
f"Try running the command manually to debug: `{' '.join(ffmpeg_cmd)}`"
)
# Verify frame count
ffprobe_path = shutil.which("ffprobe")
if not ffprobe_path:
raise OSError(
"ffprobe is required to validate the output video frame count but was not found in PATH. "
"Install ffprobe (usually distributed with ffmpeg) and retry."
)
info = get_video_info(str(output_path))
actual_fps = info["video.fps"]
result = subprocess.run(
[
ffprobe_path,
"-v",
"error",
"-select_streams",
"v:0",
"-count_frames",
"-show_entries",
"stream=nb_read_frames",
"-of",
"csv=p=0",
str(output_path),
],
capture_output=True,
text=True,
)
if result.returncode != 0 or not result.stdout.strip().isdigit():
raise OSError(
f"Failed to validate output video frame count with ffprobe for {output_path}. "
f"Return code: {result.returncode}. stderr: {result.stderr.strip()}"
)
actual_frames = int(result.stdout.strip())
if abs(actual_frames - num_frames) > 1:
raise ValueError(
f"Resampled video has {actual_frames} frames but expected {num_frames}. "
f"Source: {input_path}, target FPS: {target_fps}, duration: {duration:.4f}s. "
f"Actual FPS in output: {actual_fps}"
)