PyNvVideoCodec 2.2.1__cp313-cp313-win_arm64.whl

samples/utils/Utils.py

@@ -0,0 +1,405 @@
+# 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.
+
+import sys
+import logging
+
+from struct import pack
+import sys
+import os
+import argparse
+from pathlib import Path
+from enum import Enum
+
+import pycuda.driver as cuda
+import PyNvVideoCodec as nvc
+import numpy as np
+from pycuda.compiler import SourceModule
+import ctypes as C
+from typing import List
+import torch
+
+import pycuda.autoinit as context
+import json
+
+from contextlib import contextmanager
+
+import io
+import tempfile
+
+SERVICE_LOGGING_FORMAT = (
+        "[{filename:s}][{funcName:s}:{lineno:d}]" + "[{levelname:s}] {message:s}"
+)
+SERVICE_LOGGING_STREAM = sys.stdout
+
+
+def get_logger(logger_name, log_level="info"):
+    SERVICE_LOGGING_LEVEL = getattr(logging, log_level.upper(), None)
+
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(SERVICE_LOGGING_LEVEL)
+    ch = logging.StreamHandler(SERVICE_LOGGING_STREAM)
+    formatter = logging.Formatter(SERVICE_LOGGING_FORMAT, style="{")
+    ch.setFormatter(formatter)
+    ch.setLevel(SERVICE_LOGGING_LEVEL)
+    logger.addHandler(ch)
+    logger.propagate = False
+
+    return logger
+
+
+logger = get_logger(__file__)
+
+
+def cast_address_to_1d_bytearray(base_address, size):
+    return np.ctypeslib.as_array(C.cast(base_address, C.POINTER(C.c_uint8)),
+                                 shape=(size,))
+
+
+from contextlib import contextmanager
+import ctypes
+import io
+import os, sys
+import tempfile
+
+class AppCAI:
+    def __init__(self, shape, stride, typestr, gpualloc):
+        shape_int = tuple([int(x) for x in shape])
+        stride_int = tuple([int(x) for x in stride])
+        self.__cuda_array_interface__ = {"shape": shape_int, "strides": stride_int, "data": (int(gpualloc), False),
+                                         "typestr": typestr, "version": 3}
+
+
+class AppFrame:
+    def __init__(self, width, height, format):
+        if format == "NV12":
+            nv12_frame_size = int(width * height * 3 / 2)
+            self.gpuAlloc = cuda.mem_alloc(nv12_frame_size)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            chroma_alloc = int(self.gpuAlloc) + width * height
+            self.cai.append(AppCAI((int(height / 2), int(width / 2), 2), (width, 2, 1), "|u1", chroma_alloc))
+            self.frameSize = nv12_frame_size
+        if format == "ARGB" or format == "ABGR":
+            self.frameSize = width * height * 4
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = AppCAI((height, width, 4), (4 * width, 4, 1), "|u1", self.gpuAlloc)
+        if format == "YUV444":
+            self.frameSize = width * height * 3
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", int(self.gpuAlloc) + width * height))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", int(self.gpuAlloc) + 2 * width * height))
+        if format == "YUV420":
+            self.frameSize = int(width * height * 3 / 2)
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(
+                AppCAI((height / 2, width / 2, 1), (width / 2, 1, 1), "|u1", int(self.gpuAlloc) + width * height))
+            self.cai.append(AppCAI((height / 2, width / 2, 1), (width / 2, 1, 1), "|u1",
+                                   int(self.gpuAlloc) + width * height + width / 2 * height / 2))
+        if format == "P010" or format == "YUV420_10BIT":
+            self.frameSize = int(width * height * 3 / 2 * 2)
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            chroma_alloc = int(self.gpuAlloc) + width * height * 2
+            self.cai.append(AppCAI((int(height / 2), int(width / 2), 2), (width*2, 2, 1), "|u2", chroma_alloc))
+        if format == "P016" or format == "YUV444_16BIT":
+            self.frameSize = int(width * height * 3 * 2)
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            self.cai.append(
+                AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 2))
+            self.cai.append(
+                AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 4))
+        
+        if format == "NV16":
+            self.frameSize = width * height * 2
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1 ), "|u1", int(self.gpuAlloc) + width * height))
+
+        if format == "P210":
+            self.frameSize = width * height * 2 * 2
+            self.gpuAlloc = cuda.mem_alloc(self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 2))
+
+        
+
+    def cuda(self):
+        return self.cai
+
+
+class AppFramePerf:
+    def __init__(self, width, height, format, dataptr, frame_idx):
+        if format == "NV12":
+            nv12_frame_size = int(width * height * 3 / 2)
+            self.gpuAlloc = int(dataptr) + (frame_idx * nv12_frame_size)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            chroma_alloc = int(self.gpuAlloc) + width * height
+            self.cai.append(AppCAI((int(height / 2), int(width / 2), 2), (width, 2, 1), "|u1", chroma_alloc))
+            self.frameSize = nv12_frame_size
+        if format == "ARGB" or format == "ABGR":
+            self.frameSize = width * height * 4
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = AppCAI((height, width, 4), (4 * width, 4, 1), "|u1", self.gpuAlloc)
+        if format == "YUV444":
+            self.frameSize = width * height * 3
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", int(self.gpuAlloc) + width * height))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", int(self.gpuAlloc) + 2 * width * height))
+        if format == "YUV420":
+            self.frameSize = int(width * height * 3 / 2)
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(
+                AppCAI((height / 2, width / 2, 1), (width / 2, 1, 1), "|u1", int(self.gpuAlloc) + width * height))
+            self.cai.append(AppCAI((height / 2, width / 2, 1), (width / 2, 1, 1), "|u1",
+                                   int(self.gpuAlloc) + width * height + width / 2 * height / 2))
+        if format == "P010" or format == "YUV420_10BIT":
+            self.frameSize = int(width * height * 3 / 2 * 2)
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            chroma_alloc = int(self.gpuAlloc) + width * height * 2
+            self.cai.append(AppCAI((int(height / 2), int(width / 2), 2), (width*2, 2, 1), "|u2", chroma_alloc))
+        if format == "P016" or format == "YUV444_16BIT":
+            self.frameSize = int(width * height * 3 * 2)
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            self.cai.append(
+                AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 2))
+            self.cai.append(
+                AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 4))
+        if format == "NV16":
+            self.frameSize = width * height * 2
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width, 1, 1), "|u1", int(self.gpuAlloc) + width * height))
+        if format == "P210":
+            self.frameSize = width * height * 2 * 2
+            self.gpuAlloc = int(dataptr) + (frame_idx * self.frameSize)
+            self.cai = []
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", self.gpuAlloc))
+            self.cai.append(AppCAI((height, width, 1), (width * 2, 2, 1), "|u2", int(self.gpuAlloc) + width * height * 2))
+
+    def cuda(self):
+        return self.cai
+
+
+def FetchGPUFrame(input_frame_list, GetCPUFrameFunc, num_frames):
+    for i in range(num_frames):
+        n = i % len(input_frame_list)
+        raw_frame = GetCPUFrameFunc()
+        if not raw_frame.size:
+            return
+        cuda.memcpy_htod(input_frame_list[n].gpuAlloc, raw_frame)
+        cuda.Context.synchronize()
+        yield input_frame_list[n]
+
+
+
+def FetchCPUFrame(dec_file, frame_size):
+    def InnerFunc():
+        return np.fromfile(dec_file, np.uint8, count=frame_size)
+
+    return InnerFunc
+
+def GetChromaFormat(demuxer):
+    strFormat = "NV12"
+
+    bitdepth = demuxer.BitDepth()
+    chromaformat = demuxer.ChromaFormat()
+
+    #print (" bitdepth : " + str(bitdepth))
+    #print (" chromaformat : " + str(chromaformat))
+    
+    if str(chromaformat)  == str("cudaVideoChromaFormat.444"):
+        if bitdepth == 10:
+            strFormat = "YUV444_10BIT"
+        elif bitdepth == 16:
+            strFormat = "YUV444_16BIT"
+        elif bitdepth == 12:
+            strFormat = "YUV444_16BIT"            
+            print ("Unsupported bitdepth for YUV444 12BIT. It will be treated as YUV444_16BIT" )
+        elif bitdepth == 8:
+            strFormat = "YUV444"
+    elif str(chromaformat)  == str("cudaVideoChromaFormat.420"):
+        if bitdepth == 8:
+            strFormat = "NV12"
+        elif bitdepth == 10:
+            strFormat = "YUV420_10BIT"
+        elif bitdepth == 16:
+            strFormat = "P016"
+    elif str(chromaformat)  == str("cudaVideoChromaFormat.422"):
+        if bitdepth == 10:
+            strFormat = "YUV422_10BIT"
+        elif bitdepth == 16:
+            strFormat = "YUV422_16BIT"
+        elif bitdepth == 8:
+            strFormat = "YUV422"
+
+    #print (" strFormat : " + str(strFormat))
+
+    return strFormat
+
+
+def GetVideoCodec(demuxer):
+    codec=str(demuxer.GetNvCodecId())
+    #print (" codec : " + str(codec))
+    
+    # Common prefix to remove
+    common_prefix = "cudaVideoCodec."
+    
+    strCodec = codec.removeprefix(common_prefix)
+
+    #print (" codec : " + str(strCodec))
+    return str(strCodec)
+
+
+def get_frame_data(frame, use_device_memory=False):
+    """
+    Get raw frame data from a decoded frame.
+    
+    Args:
+        frame: DecodedFrame object from PyNvVideoCodec
+        use_device_memory: Boolean indicating if frame data is in device memory
+    
+    Returns:
+        Raw frame data (numpy array if device memory, bytearray if host memory)
+    """
+    raw_frame = np.ndarray(shape=frame.framesize(), dtype=np.uint8)
+    luma_base_addr = frame.GetPtrToPlane(0)
+    
+    if use_device_memory:
+        cuda.memcpy_dtoh(raw_frame, luma_base_addr)
+        return raw_frame
+    else:
+        return cast_address_to_1d_bytearray(base_address=luma_base_addr, size=frame.framesize())
+
+def convert_frames_to_torch(decoded_frames: List, color_format: nvc.OutputColorType = nvc.OutputColorType.RGB) -> torch.Tensor:
+    """
+    Convert batched frames from decoder to torch tensor.
+    
+    Args:
+        decoded_frames: List of decoded frames from PyNvVideoCodec
+        color_format: OutputColorType enum (RGB or RGBP)
+        
+    Returns:
+        torch.Tensor: Stacked tensor of shape (T, H, W, C) where T is number of frames
+    """
+    if not decoded_frames:
+        raise ValueError("No frames provided for conversion")
+    
+    try:
+        if color_format == nvc.OutputColorType.RGB:
+            # For RGB format, frames are already in HWC format
+            torch_frames = [torch.from_dlpack(frame) for frame in decoded_frames]
+            batched_tensor = torch.stack(torch_frames)
+            
+        elif color_format == nvc.OutputColorType.RGBP:
+            # For RGBP format, frames are in planar format
+            torch_frames = []
+            for frame in decoded_frames:
+                # Convert frame to torch tensor
+                planar_tensor = torch.from_dlpack(frame)
+                
+                # For RGBP, we get a tensor in shape (H*3, W)
+                if len(planar_tensor.shape) == 2:
+                    height = planar_tensor.shape[0] // 3
+                    width = planar_tensor.shape[1]
+                    
+                    # Split into R,G,B planes
+                    r_plane = planar_tensor[:height]
+                    g_plane = planar_tensor[height:2*height]
+                    b_plane = planar_tensor[2*height:3*height]
+                    
+                    # Stack planes along new dimension to get (H,W,3)
+                    frame_tensor = torch.stack([r_plane, g_plane, b_plane], dim=-1)
+                else:
+                    frame_tensor = planar_tensor
+                
+                torch_frames.append(frame_tensor)
+            
+            batched_tensor = torch.stack(torch_frames)
+        else:
+            raise ValueError(f"Unsupported color format: {color_format}. Only RGB and RGBP formats are supported.")
+        
+        return batched_tensor
+        
+    except Exception as e:
+        print(f"Error converting frames to torch tensor: {e}")
+        raise
+
+def frame_to_numpy(frame, color_format: nvc.OutputColorType, use_device_memory=False, width=None, height=None):
+    """
+    Convert a decoded frame to an RGB numpy array suitable for display.
+    
+    Args:
+        frame: DecodedFrame object from PyNvVideoCodec
+        color_format: nvc.OutputColorType enum (RGB or RGBP)
+        use_device_memory: Boolean indicating if frame data is in device memory
+        width: Optional frame width (if None, uses frame.width)
+        height: Optional frame height (if None, uses frame.height)
+    
+    Returns:
+        numpy.ndarray: RGB frame data as a numpy array in HWC format, ready for display
+    """
+        
+    # Get frame data using the helper function
+    frame_data = get_frame_data(frame, use_device_memory)
+    
+    # Convert to numpy array
+    frame_array = np.array(frame_data)
+    
+    # Handle different color formats
+    if color_format == nvc.OutputColorType.RGB:
+        # For RGB format, we expect 3 channels in HWC format
+        return frame_array.reshape(height, width, 3)
+    elif color_format == nvc.OutputColorType.RGBP:
+        # For planar RGB format, each plane is height*width
+        plane_size = height * width
+        if len(frame_array) != 3 * plane_size:
+            raise ValueError(f"RGBP format expects {3 * plane_size} bytes but got {len(frame_array)}")
+        r_plane = frame_array[:plane_size].reshape(height, width)
+        g_plane = frame_array[plane_size:2*plane_size].reshape(height, width)
+        b_plane = frame_array[2*plane_size:].reshape(height, width)
+        # Stack the planes to create HWC format
+        return np.dstack((r_plane, g_plane, b_plane))
+    else:
+        raise ValueError(f"Unsupported color format: {color_format}. Only RGB and RGBP formats are supported.")
+
+

samples/utils/__init__.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""
+Utility modules for PyNvVideoCodec samples.
+
+This package contains shared functionality used across multiple samples:
+- Utils.py: Common video codec utilities
+- decode_parser.py: Command line argument parsing for decoder samples
+- decode_stats_writer.py: Decode statistics extraction and analysis
+- sei_message_parser.py: SEI message extraction and parsing
+- transcode_parser.py: Transcoder command line argument parsing
+- encode_parser.py: Encoder command line argument parsing
+- frame_utils.py: Helper module for frame and memory management in video encoding/decoding.
+- encode_parallel_utils.py: Parallel utilities for encoding
+"""