nemo_curator.utils.nvcodec_utils

Module Contents

Classes

Name	Description
`FrameExtractionPolicy`	Policy for extracting frames from video, supporting full extraction or FPS-based sampling.
`NvVideoDecoder`	Low-level NVIDIA hardware-accelerated video decoder.
`PyNvcFrameExtractor`	High-level frame extraction interface using PyNvVideoCodec.
`VideoBatchDecoder`	GPU-accelerated video decoder that processes video frames in batches.

Functions

Name	Description
`gpu_decode_for_stitching`	Decode video frames for stitching using GPU acceleration.

Data

_PYNVC_AVAILABLE

pixel_format_to_cvcuda_code

API

class nemo_curator.utils.nvcodec_utils.FrameExtractionPolicy

Bases: enum.Enum

Policy for extracting frames from video, supporting full extraction or FPS-based sampling.

This enum defines the available strategies for frame extraction from video content.

fps

= 1

full

= 0

class nemo_curator.utils.nvcodec_utils.NvVideoDecoder(
    enc_file: str,
    device_id: int,
    batch_size: int,
    cuda_ctx: typing.Any,
    cvcuda_stream: typing.Any
)

Low-level NVIDIA hardware-accelerated video decoder.

This class provides direct access to NVIDIA’s hardware video decoding capabilities, handling frame decoding and memory management for video processing pipelines.

decoded_frame_cnt

= 0

input_frame_list

= queue.Queue()

local_frame_index

= 0

nvDec

nvDemux

= Nvc.PyNvDemuxer(self.input_path)

pixelFormat

= self.nvDec.GetPixelFormat()

sent_frame_cnt

= 0

nemo_curator.utils.nvcodec_utils.NvVideoDecoder.generate_decoded_frames() -> list[torch.Tensor]

Generate decoded frames from the video.

Returns: list[torch.Tensor]

List of decoded frames as tensors.

nemo_curator.utils.nvcodec_utils.NvVideoDecoder.get_next_frames() -> torch.Tensor | None

Get the next frames from the video.

Returns: torch.Tensor | None

Next frames from the video as a tensor.

class nemo_curator.utils.nvcodec_utils.PyNvcFrameExtractor(
    width: int,
    height: int,
    batch_size: int
)

High-level frame extraction interface using PyNvVideoCodec.

This class provides a simplified interface for extracting frames from videos using hardware acceleration, supporting both full extraction and FPS-based sampling.

decoder

torch_stream

= torch.cuda.ExternalStream(cvcuda_stream.handle)

nemo_curator.utils.nvcodec_utils.PyNvcFrameExtractor.__call__(
    input_path: pathlib.Path,
    extraction_policy: nemo_curator.utils.nvcodec_utils.FrameExtractionPolicy = FrameExtractionPolicy.full,
    sampling_fps: int = 1
) -> torch.Tensor

Extract frames from the video.

Parameters:

input_path

Path

Path to the video file to process.

extraction_policy

FrameExtractionPolicyDefaults to FrameExtractionPolicy.full

Policy for extracting frames.

sampling_fps

intDefaults to 1

Sampling rate for FPS-based extraction.

Returns: torch.Tensor

List of decoded frames as tensors.

class nemo_curator.utils.nvcodec_utils.VideoBatchDecoder(
    batch_size: int,
    target_width: int,
    target_height: int,
    device_id: int,
    cuda_ctx: str,
    cvcuda_stream: str
)

GPU-accelerated video decoder that processes video frames in batches.

This class handles video decoding using NVIDIA hardware acceleration, supporting batch processing of frames with color space conversion and resizing capabilities.

decoder

NvVideoDecoder | None = None

input_path

str | None = None

prev_batch_size

= self.batch_size

nemo_curator.utils.nvcodec_utils.VideoBatchDecoder.__call__(
    input_path: str
) -> torch.Tensor | None

Process video frames in batches using GPU acceleration.

Parameters:

input_path

str

Path to the video file to process.

Returns: torch.Tensor | None

Processed video frames as a tensor.

nemo_curator.utils.nvcodec_utils.VideoBatchDecoder.get_fps() -> int | None

Get the frame rate of the video.

Returns: int | None

Frame rate of the video.

nemo_curator.utils.nvcodec_utils.gpu_decode_for_stitching(
    device_id: int,
    ctx: str,
    stream: int,
    input_path: pathlib.Path,
    frame_list: list[int],
    batch_size: int = 1
) -> list[torch.Tensor]

Decode video frames for stitching using GPU acceleration.

Parameters:

device_id

int

GPU device ID.

ctx

str

CUDA context.

stream

int

CUDA stream.

input_path

Path

Path to the video file to process.

frame_list

list[int]

List of frame indices to decode.

batch_size

intDefaults to 1

Number of frames to process in each batch.

Returns: list[torch.Tensor]

List of decoded frames as tensors.

nemo_curator.utils.nvcodec_utils._PYNVC_AVAILABLE = True

nemo_curator.utils.nvcodec_utils.pixel_format_to_cvcuda_code = {Nvc.Pixel_Format.YUV444: cvcuda.ColorConversion.YUV2RGB, Nvc.Pixel_Format.NV12:...