PyNvVideoCodec 2.2.1__cp312-cp312-win_arm64.whl

samples/basic/encode.py

@@ -0,0 +1,574 @@
+# This copyright notice applies to this file only
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+'''
+This sample demonstrates unified video encoding with support for both CPU and GPU buffer modes,
+with optional muxing to container formats (MP4, MOV, WEBM).
+
+It shows how to:
+1. Configure and use different buffer modes for encoding
+2. Handle both host (CPU) and device (GPU) memory efficiently
+3. Manage encoder resources and states
+4. Process frames with optimal memory handling
+5. Mux encoded streams into container formats
+
+The sample supports two buffer modes:
+- CPU Buffer Mode:
+  * Uses host memory for frame data
+
+- GPU Buffer Mode:
+  * Uses device memory for frame data
+
+Usage:
+    python encode.py -i <input_file> -s <size> -m <mode> [options]
+
+Note:
+    FPS and GOP size are configured via the JSON config file (see -json).
+    Codec can be set via -c on the CLI (takes priority) or via the JSON config.
+    The default JSON config sets codec=h264, fps=30, gop=30.
+
+Examples:
+    # Encode to MP4 container (auto-mux from .mp4 extension)
+    python encode.py -i input.yuv -s 1920x1080 -m gpu -o output.mp4
+    
+    # Encode to raw H.264 bitstream (auto: .h264 extension = no muxing)
+    python encode.py -i input.yuv -s 1920x1080 -m gpu -o output.h264
+
+    # AV1 elementary (.av1): IVF container + patched frame count (AppEncCuda-compatible); use MP4 output for container muxing.
+    # Set "codec": "av1" in the JSON config file before running.
+    python encode.py -i input.yuv -s 1920x1080 -m cpu -json my_av1_config.json -o output.av1
+
+Arguments:
+    -i: Raw input video file
+    -s: Frame size (e.g., 1920x1080)
+    -m: Buffer mode (cpu/gpu)
+    -if: Input format (default: NV12)
+    -f: Number of frames (default: all)
+    -g: GPU ID (default: 0)
+    -c: Output codec (h264/hevc/av1); overrides codec in JSON config
+    -json: Encoder config file (controls codec, fps, gop, preset, rc, bitrate, etc.)
+    -o: Output file (use a container extension like .mp4 or .mov to enable muxing)
+'''
+
+import PyNvVideoCodec as nvc
+import numpy as np
+import json
+import argparse
+from pathlib import Path
+import os
+import sys
+from os.path import join, dirname, abspath
+
+# Add the parent directory to Python path in a way that works from any working directory
+current_dir = dirname(abspath(__file__))
+parent_dir = dirname(current_dir)
+if parent_dir not in sys.path:
+    sys.path.append(parent_dir)
+
+# Import utilities for GPU buffer mode
+try:
+    from utils.Utils import AppFrame, FetchCPUFrame, FetchGPUFrame
+    GPU_UTILS_AVAILABLE = True
+except ImportError:
+    GPU_UTILS_AVAILABLE = False
+    print("Warning: Utils module not found. GPU buffer mode will be disabled.")
+
+
+# Import utilities for frame handling
+from utils.frame_utils import get_frame_size
+from utils.encode_parser import parse_unified_args
+
+
+
+
+# Container formats that support muxing
+CONTAINER_FORMATS = {'.mp4', '.mov', '.mkv', '.webm', '.avi'}
+
+# Supported codec+container combinations for muxing.
+# Keys are extensions, values are the set of codecs that work in that container.
+_SUPPORTED_MUX = {
+    '.mp4': {'h264', 'hevc', 'av1'},
+    '.mov': {'h264', 'hevc'},
+}
+
+# Standard MPEG timebase (90 kHz) used for muxed output.
+_TIMEBASE = 90000
+
+
+def should_mux(output_path: str) -> bool:
+    """Check if output file extension indicates a container format."""
+    ext = Path(output_path).suffix.lower()
+    return ext in CONTAINER_FORMATS
+
+
+def validate_mux_support(codec: str, output_path: str) -> None:
+    """
+    Raise ValueError if the codec+container combination is not supported.
+
+    Supported combinations:
+        MP4: H264, HEVC, AV1
+        MOV: H264, HEVC
+    """
+    ext = Path(output_path).suffix.lower()
+    supported_codecs = _SUPPORTED_MUX.get(ext)
+    if supported_codecs is None:
+        supported = ', '.join(sorted(f'*{e}' for e in _SUPPORTED_MUX))
+        raise ValueError(
+            f"Muxing to '{ext}' is not supported. Supported containers: {supported}"
+        )
+    if codec.lower() not in supported_codecs:
+        raise ValueError(
+            f"Muxing {codec.upper()} into '{ext}' is not supported. "
+            f"Supported codecs for {ext}: {', '.join(sorted(c.upper() for c in supported_codecs))}"
+        )
+
+
+def get_media_format(output_path: str):
+    """
+    Get media format enum from file extension.
+    
+    Uses PyNvVideoCodec's built-in GetMediaFormat function to get the
+    proper MEDIA_FORMAT enum value required by FFmpegMuxer.
+    
+    Args:
+        output_path: Path to output file
+        
+    Returns:
+        MEDIA_FORMAT enum value (e.g., nvc.MEDIA_FORMAT.MP4)
+    """
+    return nvc.GetMediaFormat(output_path)
+
+
+def _concat_encode_packets(packets) -> bytes:
+    """Join Encode / EndEncode output into one bitstream blob."""
+    return b"".join(bytes(p["data"]) for p in packets)
+
+
+
+def _av1_ivf_elementary_enabled(config_params: dict, use_muxer: bool) -> bool:
+    """True when output is AV1 in IVF (elementary), matching VideoCodecSDK AppEncCuda."""
+    if use_muxer or config_params.get("codec", "h264").lower() != "av1":
+        return False
+    v = str(config_params.get("use_ivf_container", "1")).lower()
+    return v not in ("0", "false", "no")
+
+
+def _patch_ivf_output_file(output_file_path: str, ivf_packet_count: int) -> None:
+    """Match FFmpeg ivf_write_trailer: frame count at IVF global header offset 24 (LE32)."""
+    if ivf_packet_count <= 0:
+        return
+    with open(output_file_path, "r+b") as out_f:
+        if out_f.read(4) != b"DKIF":
+            return
+        out_f.seek(24)
+        out_f.write(ivf_packet_count.to_bytes(4, "little"))
+
+
+def get_extradata_from_encoder(encoder) -> bytes:
+    """
+    Get sequence parameters (SPS/PPS/VPS) from encoder.
+    
+    This uses the NVENC API's GetSequenceParams function to retrieve the
+    codec-specific header data required by container formats.
+    
+    Args:
+        encoder: The PyNvEncoder instance
+        
+    Returns:
+        bytes: Sequence parameter data for muxer extradata
+    """
+    return encoder.GetSequenceParams()
+
+
+def encode_cpu_buffer(gpu_id, dec_file_path, enc_file_path, width, height, fmt, config_params, frame_count, gop_size, fps, use_mux=None):
+    """
+    Encode frames using host memory buffers as input.
+
+    This function reads image data from a file and copies it to CUDA buffers for encoding.
+    The encoder submits the data to NVENC hardware for encoding. Video memory buffer is
+    allocated to get the NVENC hardware output. The output is copied from video memory
+    to host memory for file storage.
+
+    Parameters:
+        gpu_id (int): Ordinal of GPU to use
+        dec_file_path (str): Path to file to be decoded
+        enc_file_path (str): Path to output file for encoded frames
+        width (int): Width of encoded frame
+        height (int): Height of encoded frame
+        fmt (str): Surface format string in uppercase (e.g., "NV12")
+        config_params (dict): Key-value pairs providing fine-grained control on encoding
+        frame_count (int): Number of frames to encode
+        gop_size (int): GOP size (keyframe interval)
+        fps (int): Frame rate (frames per second) for muxing
+        use_mux (bool): If True, mux to container. If False, raw bitstream. If None, auto-detect.
+
+    Returns:
+        None
+    """
+    print("Using CPU buffer mode")
+    
+    # Check if we need to mux (use override if provided, otherwise auto-detect)
+    use_muxer = use_mux if use_mux is not None else should_mux(enc_file_path)
+    codec = config_params.get("codec", "h264").lower()
+    if codec == "av1":
+        if use_muxer:
+            config_params["use_ivf_container"] = "0"
+        else:
+            config_params.setdefault("use_ivf_container", "1")
+
+    frame_size = get_frame_size(width, height, fmt)
+    
+    if use_muxer:
+        validate_mux_support(codec, enc_file_path)
+        print(f"Will mux output to container format: {Path(enc_file_path).suffix}")
+
+    with open(dec_file_path, "rb") as dec_file:
+        config_params["gpu_id"] = gpu_id
+        nvenc = nvc.CreateEncoder(width, height, fmt, True, **config_params)  # True = use CPU buffers
+        
+        # Calculate total available frames from file size
+        file_size = os.path.getsize(dec_file_path)
+        available_frames = file_size // frame_size
+        
+        # Determine frames to process with limit
+        frames_to_process = frame_count if frame_count > 0 else available_frames
+        frames_to_process = min(frames_to_process, available_frames)
+        
+        print(f"File contains {available_frames} frames, processing {frames_to_process} frames using CPU buffers...")
+        
+        # Muxing parameters
+        pts_increment = _TIMEBASE // fps
+
+        # Initialize muxer or output file
+        muxer = None
+        enc_file = None
+
+        if use_muxer:
+            # Get sequence parameters (SPS/PPS) from encoder for muxer extradata
+            extradata = get_extradata_from_encoder(nvenc)
+            print(f"Extracted extradata size: {len(extradata)} bytes")
+
+            media_format = get_media_format(enc_file_path)
+            muxer = nvc.FFmpegMuxer(
+                file_path=enc_file_path,
+                media_format=media_format,
+                codec=codec,
+                width=width,
+                height=height,
+                fps_num=fps,
+                fps_den=1,
+                timebase_num=1,
+                timebase_den=_TIMEBASE,
+                extradata=extradata
+            )
+            # Inform the muxer of the uniform frame duration so Finalize()
+            # generates correct PTS/DTS without needing source container PTS.
+            muxer.SetUniformPtsIncrement(pts_increment)
+        else:
+            # Open raw output file
+            enc_file = open(enc_file_path, "wb")
+
+        frames_encoded = 0
+        ivf_packets = 0
+
+        for i in range(frames_to_process):
+            chunk = np.fromfile(dec_file, np.uint8, count=frame_size)
+            if chunk.size != 0:
+                pic_params = nvc.NV_ENC_PIC_PARAMS()
+                pic_params.inputTimeStamp = frames_encoded
+                enc_list = nvenc.Encode(chunk, pic_params)
+                ivf_packets += len(enc_list)
+
+                if use_muxer:
+                    for enc_packet in enc_list:
+                        muxer.MuxVideoPacket(bytes(enc_packet["data"]),
+                                             enc_packet["picture_type"],
+                                             enc_packet["timestamp"])
+                else:
+                    enc_file.write(bytearray(_concat_encode_packets(enc_list)))
+
+                frames_encoded += 1
+            else:
+                print(f"Warning: Could not read frame {i+1}")
+                break
+
+        flush_list = nvenc.EndEncode()
+        ivf_packets += len(flush_list)
+        if use_muxer:
+            for flush_packet in flush_list:
+                muxer.MuxVideoPacket(bytes(flush_packet["data"]),
+                                     flush_packet["picture_type"],
+                                     flush_packet["timestamp"])
+            muxer.Finalize()
+        else:
+            bitstream = _concat_encode_packets(flush_list)
+            enc_file.write(bytearray(bitstream))
+            if _av1_ivf_elementary_enabled(config_params, use_muxer):
+                enc_file.flush()
+                _patch_ivf_output_file(enc_file_path, ivf_packets)
+            enc_file.close()
+
+        print(f"Completed encoding {frames_encoded} frames using CPU buffers")
+        print(f"Output file: {enc_file_path}")
+
+
+def encode_gpu_buffer(gpu_id, dec_file_path, enc_file_path, width, height, fmt, config_params, frame_count, gop_size, fps, use_mux=None):
+    """
+    Encode frames using CUDA device buffers as input.
+
+    This function reads image data from a file and loads it to CUDA input buffers using
+    FetchGPUFrame(). The encoder copies the CUDA buffers and submits them to NVENC hardware
+    for encoding. Video memory buffer is allocated to get the NVENC hardware output.
+    The output is copied from video memory to host memory for file storage.
+
+    Parameters:
+        gpu_id (int): Ordinal of GPU to use
+        dec_file_path (str): Path to file to be decoded
+        enc_file_path (str): Path to output file for encoded frames
+        width (int): Width of encoded frame
+        height (int): Height of encoded frame
+        fmt (str): Surface format string in uppercase (e.g., "NV12")
+        config_params (dict): Key-value pairs providing fine-grained control on encoding
+        frame_count (int): Number of frames to encode
+        gop_size (int): GOP size (keyframe interval)
+        fps (int): Frame rate (frames per second) for muxing
+        use_mux (bool): If True, mux to container. If False, raw bitstream. If None, auto-detect.
+
+    Returns:
+        None
+    """
+    print("Using GPU buffer mode")
+    
+    if not GPU_UTILS_AVAILABLE:
+        raise RuntimeError("GPU buffer mode requires Utils module (AppFrame, FetchCPUFrame, FetchGPUFrame)")
+    
+    # Check if we need to mux (use override if provided, otherwise auto-detect)
+    use_muxer = use_mux if use_mux is not None else should_mux(enc_file_path)
+    codec = config_params.get("codec", "h264").lower()
+    if codec == "av1":
+        if use_muxer:
+            config_params["use_ivf_container"] = "0"
+        else:
+            config_params.setdefault("use_ivf_container", "1")
+
+    if use_muxer:
+        validate_mux_support(codec, enc_file_path)
+        print(f"Will mux output to container format: {Path(enc_file_path).suffix}")
+
+    # Show encoder capabilities
+    caps = nvc.GetEncoderCaps(codec=config_params.get("codec", "h264"))
+    if "num_encoder_engines" in caps:
+        print(f"Number of NVENCs: {caps['num_encoder_engines']}")
+
+    with open(dec_file_path, "rb") as decFile:
+        config_params["gpu_id"] = gpu_id
+        nvenc = nvc.CreateEncoder(width, height, fmt, False, **config_params)  # False = use GPU buffers
+        input_frame_list = [AppFrame(width, height, fmt) for _ in range(1, 5)]
+        
+        # Calculate total available frames from file size
+        file_size = os.path.getsize(dec_file_path)
+        available_frames = file_size // input_frame_list[0].frameSize
+        
+        # Determine frames to process with limit
+        frames_to_process = frame_count if frame_count > 0 else available_frames
+        frames_to_process = min(frames_to_process, available_frames)
+        
+        print(f"File contains {available_frames} frames, processing {frames_to_process} frames using GPU buffers...")
+        
+        # Muxing parameters
+        pts_increment = _TIMEBASE // fps
+
+        # Initialize muxer or output file
+        muxer = None
+        encFile = None
+
+        if use_muxer:
+            # Get sequence parameters (SPS/PPS) from encoder for muxer extradata
+            extradata = get_extradata_from_encoder(nvenc)
+            print(f"Extracted extradata size: {len(extradata)} bytes")
+
+            media_format = get_media_format(enc_file_path)
+            muxer = nvc.FFmpegMuxer(
+                file_path=enc_file_path,
+                media_format=media_format,
+                codec=codec,
+                width=width,
+                height=height,
+                fps_num=fps,
+                fps_den=1,
+                timebase_num=1,
+                timebase_den=_TIMEBASE,
+                extradata=extradata
+            )
+            # Inform the muxer of the uniform frame duration so Finalize()
+            # generates correct PTS/DTS without needing source container PTS.
+            muxer.SetUniformPtsIncrement(pts_increment)
+        else:
+            # Open raw output file
+            encFile = open(enc_file_path, "wb")
+
+        frames_encoded = 0
+        ivf_packets = 0
+
+        for input_gpu_frame in FetchGPUFrame(
+            input_frame_list,
+            FetchCPUFrame(decFile, input_frame_list[0].frameSize),
+            frames_to_process
+        ):
+            pic_params = nvc.NV_ENC_PIC_PARAMS()
+            pic_params.inputTimeStamp = frames_encoded
+            enc_list = nvenc.Encode(input_gpu_frame, pic_params)
+            ivf_packets += len(enc_list)
+
+            if use_muxer:
+                for enc_packet in enc_list:
+                    muxer.MuxVideoPacket(bytes(enc_packet["data"]),
+                                         enc_packet["picture_type"],
+                                         enc_packet["timestamp"])
+            else:
+                encFile.write(bytearray(_concat_encode_packets(enc_list)))
+
+            frames_encoded += 1
+
+        flush_list = nvenc.EndEncode()
+        ivf_packets += len(flush_list)
+        if use_muxer:
+            for flush_packet in flush_list:
+                muxer.MuxVideoPacket(bytes(flush_packet["data"]),
+                                     flush_packet["picture_type"],
+                                     flush_packet["timestamp"])
+            muxer.Finalize()
+        else:
+            bitstream = bytearray(_concat_encode_packets(flush_list))
+            encFile.write(bitstream)
+            if _av1_ivf_elementary_enabled(config_params, use_muxer):
+                encFile.flush()
+                _patch_ivf_output_file(enc_file_path, ivf_packets)
+            encFile.close()
+
+    print(f"Completed encoding {frames_to_process} frames using GPU buffers")
+    print(f"Output file: {enc_file_path}")
+
+
+def encode_unified(gpu_id, dec_file_path, enc_file_path, width, height, fmt, config_params, frame_count, buffer_mode, gop_size, fps, use_mux=None):
+    """
+    Unified encoding function that supports both CPU and GPU buffer modes.
+
+    Parameters:
+        gpu_id (int): Ordinal of GPU to use
+        dec_file_path (str): Path to file to be decoded
+        enc_file_path (str): Path to output file for encoded frames
+        width (int): Width of encoded frame
+        height (int): Height of encoded frame
+        fmt (str): Surface format string in uppercase (e.g., "NV12")
+        config_params (dict): Key-value pairs providing fine-grained control on encoding
+        frame_count (int): Number of frames to encode
+        buffer_mode (str): Buffer mode - "cpu" or "gpu"
+        gop_size (int): GOP size (keyframe interval)
+        fps (int): Frame rate (frames per second) for muxing
+        use_mux (bool): If True, mux to container. If False, raw bitstream. If None, auto-detect.
+
+    Returns:
+        None
+    """
+    try:
+        if buffer_mode.lower() == "cpu":
+            encode_cpu_buffer(gpu_id, dec_file_path, enc_file_path, width, height, fmt, config_params, frame_count, gop_size, fps, use_mux)
+        elif buffer_mode.lower() == "gpu":
+            encode_gpu_buffer(gpu_id, dec_file_path, enc_file_path, width, height, fmt, config_params, frame_count, gop_size, fps, use_mux)
+        else:
+            raise ValueError(f"Invalid buffer mode: {buffer_mode}. Must be 'cpu' or 'gpu'")
+            
+    except nvc.PyNvVCExceptionUnsupported as e:
+        print(f"CreateEncoder failure: {e}")
+    except ValueError as e:
+        print(f"Error: {e}")
+        sys.exit(1)
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+
+def validate_gpu_mode(args) -> None:
+    """Validate GPU buffer mode requirements."""
+    if args.buffer_mode == "gpu" and not GPU_UTILS_AVAILABLE:
+        print("Error: GPU buffer mode requires Utils module, but it's not available.")
+        print("Please use CPU buffer mode (-m cpu) or ensure Utils module is accessible.")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    args = parse_unified_args(
+        "Unified encoding application supporting both CPU and GPU buffer modes with optional muxing."
+    )
+
+    # GOP size and FPS are sourced from the JSON config (via load_config)
+    gop_size = args.gop_size
+    fps = args.fps
+    
+    use_mux = should_mux(str(args.encoded_file_path))
+
+    # Display configuration
+    print("=" * 60)
+    print("ENCODER CONFIGURATION")
+    print("=" * 60)
+    print(f"Buffer mode: {args.buffer_mode.upper()}")
+    print(f"Input file: {args.raw_file_path}")
+    print(f"Output file: {args.encoded_file_path}")
+    print(f"Resolution: {args.width}x{args.height}")
+    print(f"Format: {args.format}")
+    print(f"Codec: {args.codec}")
+    print(f"GPU ID: {args.gpu_id}")
+    print(f"GOP size: {gop_size}")
+    print(f"Frame rate: {fps} fps")
+    print(f"Frame count: {args.frame_count if args.frame_count > 0 else 'All frames'}")
+    
+    # Show muxing status
+    if use_mux:
+        print(f"Muxing: Enabled ({Path(args.encoded_file_path).suffix.upper()})")
+    else:
+        print(f"Muxing: Disabled (raw bitstream)")
+    
+    print("=" * 60)
+
+    # Validate GPU mode requirements
+    validate_gpu_mode(args)
+    
+    # Run encoding
+    encode_unified(
+        args.gpu_id,
+        args.raw_file_path.as_posix(),
+        args.encoded_file_path.as_posix(),
+        args.width,
+        args.height,
+        args.format,
+        args.config,
+        args.frame_count,
+        args.buffer_mode,
+        gop_size,
+        fps,
+        use_mux
+    )
+
+    print("=" * 60)
+    print("ENCODING COMPLETE")
+    print("=" * 60)

samples/basic/segments.txt

@@ -0,0 +1,3 @@
+0.3 0.8
+2 4
+3 8