Isaac-GR00T/gr00t/utils/video_utils.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

import math
from typing import List, Optional, Tuple

import numpy as np


_TORCHCODEC_INSTALL_HINT = (
    "torchcodec is required for video decoding. Install it via the platform "
    "deployment script (scripts/deployment/{dgpu,orin,thor,spark}/install_deps.sh) "
    "or `uv pip install torchcodec`."
)

_TORCHCODEC_FFMPEG_HINT = (
    "torchcodec is installed but its native library could not be loaded. This is "
    "almost always an FFmpeg version mismatch: torchcodec 0.8.0 only ships loaders "
    "for FFmpeg 4-7 and cannot run against FFmpeg 8 (the default on Ubuntu 25.10+ "
    "and 26.04). Install an FFmpeg<8 runtime (e.g. `conda install -c conda-forge "
    "'ffmpeg<8'` or a distro package) and ensure its shared libraries are on "
    "LD_LIBRARY_PATH. Original error: {error}"
)

_DEFAULT_DECODER_KWARGS = {
    "device": "cpu",
    "dimension_order": "NHWC",
    "num_ffmpeg_threads": 0,
}


def _decoder_import_error(exc: BaseException) -> ImportError:
    """Map a torchcodec import/load failure to an actionable ImportError.

    A plain ``ImportError`` means torchcodec is not installed. An ``OSError`` or
    ``RuntimeError`` means torchcodec is present but its native library failed to
    load, which is most commonly an FFmpeg major-version mismatch.
    """
    if isinstance(exc, ImportError):
        return ImportError(_TORCHCODEC_INSTALL_HINT)
    return ImportError(_TORCHCODEC_FFMPEG_HINT.format(error=exc))


def _get_video_decoder_cls():
    try:
        from torchcodec.decoders import VideoDecoder
    except (ImportError, OSError, RuntimeError) as exc:
        raise _decoder_import_error(exc) from exc
    return VideoDecoder


def _build_decoder(video_path: str, decoder_kwargs: Optional[dict]):
    video_decoder_cls = _get_video_decoder_cls()
    kwargs = {**_DEFAULT_DECODER_KWARGS, **(decoder_kwargs or {})}
    return video_decoder_cls(video_path, **kwargs)


def get_frames_by_indices(
    video_path: str,
    indices: list[int] | np.ndarray,
    decoder_kwargs: Optional[dict] = None,
) -> np.ndarray:
    decoder = _build_decoder(video_path, decoder_kwargs)
    return decoder.get_frames_at(indices=indices).data.numpy()


def get_frames_by_timestamps(
    video_path: str,
    timestamps: list[float] | np.ndarray,
    decoder_kwargs: Optional[dict] = None,
) -> np.ndarray:
    """Get frames from a video at specified timestamps.

    Args:
        video_path (str): Path to the video file.
        timestamps (list[float] | np.ndarray): Timestamps to retrieve frames for, in seconds.

    Returns:
        np.ndarray: Frames at the specified timestamps.
    """
    decoder = _build_decoder(video_path, decoder_kwargs)

    # https://docs.pytorch.org/torchcodec/stable/generated/torchcodec.decoders.VideoStreamMetadata.html#torchcodec.decoders.VideoStreamMetadata
    fps = decoder.metadata.average_fps
    interval = 1 / fps
    timestamps = np.array(timestamps).astype(np.float64)

    # Correct float precision issues in timestamps
    # E.g. for 5fps video: [1.0, 1.20000005, 1.39999998] -> [1.0, 1.2, 1.4]
    # Without this, torchcodec will read the delayed frame (e.g. 1.39999998 -> 1.2)
    # Round to nearest frame interval to prevent torchcodec from reading wrong frames.
    # Allow max 1% error from expected interval.
    closest_timestamps = np.round(timestamps / interval) * interval
    timestamp_errors = np.abs(closest_timestamps - timestamps) / interval
    invalid_mask = timestamp_errors >= 0.01
    if np.any(invalid_mask):
        invalid_indices = np.where(invalid_mask)[0]
        invalid_timestamps = timestamps[invalid_indices]
        raise ValueError(
            f"Try to read invalid timestamps {invalid_timestamps} from video {video_path} (FPS: {fps})"
        )

    timestamps = closest_timestamps
    return decoder.get_frames_played_at(seconds=timestamps).data.numpy()


def get_all_frames(
    video_path: str,
    decoder_kwargs: Optional[dict] = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Get all frames from a video.

    Returns:
        tuple[np.ndarray, np.ndarray]: Frames and timestamps.
    """
    decoder = _build_decoder(video_path, decoder_kwargs)
    frames = decoder.get_frames_at(indices=range(len(decoder)))
    return frames.data.numpy(), frames.pts_seconds.numpy()


def get_accumulate_timestamp_idxs(
    timestamps: List[float],
    start_time: float,
    dt: float,
    eps: float = 1e-5,
    next_global_idx: Optional[int] = 0,
    allow_negative=False,
) -> Tuple[List[int], List[int], int]:
    """
    For each dt window, choose the first timestamp in the window.
    Assumes timestamps sorted. One timestamp might be chosen multiple times due to dropped frames.
    next_global_idx should start at 0 normally, and then use the returned next_global_idx.
    However, when overwiting previous values are desired, set last_global_idx to None.

    Returns:
    local_idxs: which index in the given timestamps array to chose from
    global_idxs: the global index of each chosen timestamp
    next_global_idx: used for next call.
    """
    local_idxs = list()
    global_idxs = list()
    for local_idx, ts in enumerate(timestamps):
        # add eps * dt to timestamps so that when ts == start_time + k * dt
        # is always recorded as kth element (avoiding floating point errors)
        global_idx = math.floor((ts - start_time) / dt + eps)
        if (not allow_negative) and (global_idx < 0):
            continue
        if next_global_idx is None:
            next_global_idx = global_idx

        n_repeats = max(0, global_idx - next_global_idx + 1)
        for i in range(n_repeats):
            local_idxs.append(local_idx)
            global_idxs.append(next_global_idx + i)
        next_global_idx += n_repeats
    return local_idxs, global_idxs, next_global_idx