Add files using upload-large-folder tool

b6a9d87 verified 2 months ago

6.48 kB

# 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

	# 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