opentau.datasets.video_utils

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']}")

Functions

decode_video_frames(video_path, timestamps, ...)

Decodes video frames using the specified backend.

decode_video_frames_torchcodec(video_path, ...)

Loads frames associated with the requested timestamps of a video using torchcodec.

decode_video_frames_torchvision(video_path, ...)

Loads frames associated to the requested timestamps of a video

encode_video_frames(imgs_dir, video_path, fps)

Encode a sequence of images into a video file using ffmpeg.

get_audio_info(video_path)

Extract audio stream information from a video file using ffprobe.

get_image_pixel_channels(image)

Determine the number of pixel channels from a PIL Image mode.

get_safe_default_codec()

Get the default video codec backend.

get_safe_encoding_vcodec()

Return an ffmpeg encoding codec that is available on this system.

get_video_info(video_path)

Extract video stream information from a video file using ffprobe.

get_video_pixel_channels(pix_fmt)

Determine the number of pixel channels from a pixel format string.

resample_and_trim_video(input_path, ...[, ...])

Resample a video to the target FPS and trim it to exactly num_frames frames.

Classes

VideoFrame()

Provides a type for a dataset containing video frames.
class opentau.datasets.video_utils.VideoFrame[source]

Bases: object

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)) `

__init__() → None
pa_type: ClassVar[Any] = StructType(struct<path: string, timestamp: float>)
opentau.datasets.video_utils.decode_video_frames(video_path: Path | str, timestamps: list[float], tolerance_s: float, backend: str | None = None) → Tensor[source]

Decodes video frames using the specified backend.

Parameters:

  • 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:

Decoded frames.

Return type:

torch.Tensor

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.

opentau.datasets.video_utils.decode_video_frames_torchcodec(video_path: Path | str, timestamps: list[float], tolerance_s: float, device: str = 'cpu', log_loaded_timestamps: bool = False) → Tensor[source]

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.

opentau.datasets.video_utils.decode_video_frames_torchvision(video_path: Path | str, timestamps: list[float], tolerance_s: float, backend: str = 'pyav', log_loaded_timestamps: bool = False) → Tensor[source]

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.

opentau.datasets.video_utils.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[source]

Encode a sequence of images into a video file using ffmpeg.

Parameters:

  • 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

opentau.datasets.video_utils.get_audio_info(video_path: Path | str) → dict[source]

Extract audio stream information from a video file using ffprobe.

Parameters:

video_path – Path to the video file.

Returns:

  • 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).

Return type:

Dictionary containing audio information

Raises:

RuntimeError – If ffprobe command fails.

opentau.datasets.video_utils.get_image_pixel_channels(image: <module 'PIL.Image' from '/home/docs/checkouts/readthedocs.org/user_builds/opentau/envs/latest/lib/python3.10/site-packages/PIL/Image.py'>) → int[source]

Determine the number of pixel channels from a PIL Image mode.

Parameters:

image – PIL Image object.

Returns:

Number of channels (1, 2, 3, or 4).

Raises:

ValueError – If image mode is unknown.

opentau.datasets.video_utils.get_safe_default_codec() → str[source]

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.

opentau.datasets.video_utils.get_safe_encoding_vcodec() → str[source]

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”.

opentau.datasets.video_utils.get_video_info(video_path: Path | str) → dict[source]

Extract video stream information from a video file using ffprobe.

Parameters:

video_path – Path to the video file.

Returns:

  • 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().

Return type:

Dictionary containing video and audio information

Raises:

RuntimeError – If ffprobe command fails.

opentau.datasets.video_utils.get_video_pixel_channels(pix_fmt: str) → int[source]

Determine the number of pixel channels from a pixel format string.

Parameters:

pix_fmt – Pixel format string (e.g., “yuv420p”, “rgb24”).

Returns:

Number of channels (1, 3, or 4).

Raises:

ValueError – If pixel format is unknown.

opentau.datasets.video_utils.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[source]

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 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).

Parameters:

  • 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).