pixtreme 0.5.1__py3-none-any.whl

pixtreme/__init__.py

@@ -0,0 +1,198 @@
+import cupy as cp
+
+from .color import (
+    aces2065_1_to_acescct,
+    aces2065_1_to_acescg,
+    aces2065_1_to_rec709,
+    acescct_to_aces2065_1,
+    acescg_to_aces2065_1,
+    apply_lut,
+    apply_lut_cp,
+    bgr_to_grayscale,
+    bgr_to_hsv,
+    bgr_to_rgb,
+    bgr_to_ycbcr,
+    hsv_to_bgr,
+    hsv_to_rgb,
+    ndi_uyvy422_to_ycbcr444,
+    ndi_uyvy422_to_ycbcr444_cp,
+    read_lut,
+    rec709_to_aces2065_1,
+    rgb_to_bgr,
+    rgb_to_grayscale,
+    rgb_to_hsv,
+    rgb_to_ycbcr,
+    uyvy422_to_ycbcr444,
+    uyvy422_to_ycbcr444_cp,
+    ycbcr_full_to_legal,
+    ycbcr_legal_to_full,
+    ycbcr_to_bgr,
+    ycbcr_to_grayscale,
+    ycbcr_to_rgb,
+    yuv420p_to_ycbcr444,
+    yuv420p_to_ycbcr444_cp,
+    yuv422p10le_to_ycbcr444,
+    yuv422p10le_to_ycbcr444_cp,
+)
+from .draw import add_label, circle, create_rounded_mask, put_text, rectangle
+from .filter import GaussianBlur, gaussian_blur, get_gaussian_kernel
+from .io import destroy_all_windows, imdecode, imencode, imread, imshow, imwrite, waitkey
+from .transform import (
+    INTER_AREA,
+    INTER_AUTO,
+    INTER_B_SPLINE,
+    INTER_CATMULL_ROM,
+    INTER_CUBIC,
+    INTER_LANCZOS2,
+    INTER_LANCZOS3,
+    INTER_LANCZOS4,
+    INTER_LINEAR,
+    INTER_MITCHELL,
+    INTER_NEAREST,
+    add_padding,
+    affine_transform,
+    create_erode_kernel,
+    create_gaussian_weights,
+    erode,
+    get_inverse_matrix,
+    merge_tiles,
+    resize,
+    stack_images,
+    subsample_image,
+    subsample_image_back,
+    tile_image,
+)
+from .upscale import OnnxUpscaler, TorchUpscaler, TrtUpscaler
+from .utils import (
+    batch_to_images,
+    check_onnx_model,
+    check_torch_model,
+    guess_image_layout,
+    image_to_batch,
+    images_to_batch,
+    onnx_to_onnx_dynamic,
+    onnx_to_trt,
+    onnx_to_trt_dynamic_shape,
+    onnx_to_trt_fixed_shape,
+    to_cupy,
+    to_dtype,
+    to_float16,
+    to_float32,
+    to_float64,
+    to_numpy,
+    to_tensor,
+    to_uint8,
+    to_uint16,
+    torch_to_onnx,
+)
+
+get_device_id = cp.cuda.device.get_device_id
+get_device_count = cp.cuda.runtime.getDeviceCount
+Device = cp.cuda.Device
+
+__all__ = [
+    # Color conversion functions
+    "aces2065_1_to_acescct",
+    "aces2065_1_to_acescg",
+    "aces2065_1_to_rec709",
+    "acescct_to_aces2065_1",
+    "acescg_to_aces2065_1",
+    "apply_lut",
+    "apply_lut_cp",
+    "bgr_to_grayscale",
+    "bgr_to_hsv",
+    "bgr_to_rgb",
+    "bgr_to_ycbcr",
+    "hsv_to_bgr",
+    "hsv_to_rgb",
+    "ndi_uyvy422_to_ycbcr444",
+    "ndi_uyvy422_to_ycbcr444_cp",
+    "read_lut",
+    "rec709_to_aces2065_1",
+    "rgb_to_bgr",
+    "rgb_to_grayscale",
+    "rgb_to_hsv",
+    "rgb_to_ycbcr",
+    "uyvy422_to_ycbcr444",
+    "uyvy422_to_ycbcr444_cp",
+    "ycbcr_full_to_legal",
+    "ycbcr_legal_to_full",
+    "ycbcr_to_bgr",
+    "ycbcr_to_grayscale",
+    "ycbcr_to_rgb",
+    "yuv420p_to_ycbcr444",
+    "yuv420p_to_ycbcr444_cp",
+    "yuv422p10le_to_ycbcr444",
+    "yuv422p10le_to_ycbcr444_cp",
+    # Drawing functions
+    "add_label",
+    "circle",
+    "create_rounded_mask",
+    "put_text",
+    "rectangle",
+    # Filtering
+    "GaussianBlur",
+    "gaussian_blur",
+    "get_gaussian_kernel",
+    # I/O operations
+    "destroy_all_windows",
+    "imdecode",
+    "imencode",
+    "imread",
+    "imshow",
+    "imwrite",
+    "waitkey",
+    # Transform operations
+    "INTER_AREA",
+    "INTER_AUTO",
+    "INTER_B_SPLINE",
+    "INTER_CATMULL_ROM",
+    "INTER_CUBIC",
+    "INTER_LANCZOS2",
+    "INTER_LANCZOS3",
+    "INTER_LANCZOS4",
+    "INTER_LINEAR",
+    "INTER_MITCHELL",
+    "INTER_NEAREST",
+    "affine_transform",
+    "create_erode_kernel",
+    "erode",
+    "get_inverse_matrix",
+    "resize",
+    "stack_images",
+    "subsample_image",
+    "subsample_image_back",
+    # Upscaling
+    "OnnxUpscaler",
+    "TorchUpscaler",
+    "TrtUpscaler",
+    "add_padding",
+    "create_gaussian_weights",
+    "merge_tiles",
+    "tile_image",
+    # Utilities
+    "check_onnx_model",
+    "check_torch_model",
+    "torch_to_onnx",
+    "onnx_to_onnx_dynamic",
+    "onnx_to_trt",
+    "onnx_to_trt_dynamic_shape",
+    "onnx_to_trt_fixed_shape",
+    "batch_to_images",
+    "image_to_batch",
+    "images_to_batch",
+    "guess_image_layout",
+    "to_cupy",
+    "to_dtype",
+    "to_float16",
+    "to_float32",
+    "to_float64",
+    "to_numpy",
+    "to_tensor",
+    "to_uint8",
+    "to_uint16",
+    # Device functions
+    "Device",
+    "get_device_count",
+    "get_device_id",
+]

pixtreme/color/__init__.py

@@ -0,0 +1,62 @@
+from .aces import (
+    aces2065_1_to_acescct,
+    aces2065_1_to_acescg,
+    aces2065_1_to_rec709,
+    acescct_to_aces2065_1,
+    acescg_to_aces2065_1,
+    rec709_to_aces2065_1,
+)
+from .bgr import bgr_to_rgb, rgb_to_bgr
+from .grayscale import bgr_to_grayscale, rgb_to_grayscale
+from .hsv import bgr_to_hsv, hsv_to_bgr, hsv_to_rgb, rgb_to_hsv
+from .lut import apply_lut, read_lut
+from .uyvy422 import uyvy422_to_ycbcr444
+from .uyvy422_ndi import ndi_uyvy422_to_ycbcr444
+from .ycbcr import (
+    bgr_to_ycbcr,
+    rgb_to_ycbcr,
+    ycbcr_full_to_legal,
+    ycbcr_legal_to_full,
+    ycbcr_to_bgr,
+    ycbcr_to_grayscale,
+    ycbcr_to_rgb,
+)
+from .yuv420 import yuv420p_to_ycbcr444
+from .yuv422p10le import yuv422p10le_to_ycbcr444
+
+__all__ = [
+    # ACES color management
+    "aces2065_1_to_acescct",
+    "aces2065_1_to_acescg",
+    "aces2065_1_to_rec709",
+    "acescct_to_aces2065_1",
+    "acescg_to_aces2065_1",
+    "rec709_to_aces2065_1",
+    # 3D LUT operations
+    "apply_lut",
+    "read_lut",
+    # BGR/RGB conversions
+    "bgr_to_rgb",
+    "rgb_to_bgr",
+    # Grayscale conversions
+    "bgr_to_grayscale",
+    "rgb_to_grayscale",
+    # HSV conversions
+    "bgr_to_hsv",
+    "hsv_to_bgr",
+    "hsv_to_rgb",
+    "rgb_to_hsv",
+    # YCbCr conversions
+    "bgr_to_ycbcr",
+    "rgb_to_ycbcr",
+    "ycbcr_full_to_legal",
+    "ycbcr_legal_to_full",
+    "ycbcr_to_bgr",
+    "ycbcr_to_grayscale",
+    "ycbcr_to_rgb",
+    # Video format conversions
+    "uyvy422_to_ycbcr444",
+    "ndi_uyvy422_to_ycbcr444",
+    "yuv420p_to_ycbcr444",
+    "yuv422p10le_to_ycbcr444",
+]

pixtreme/color/aces/__init__.py

@@ -0,0 +1,67 @@
+"""
+ACES Color Management
+
+Provides ACES 1.2 color transformations following Nuke/VFX industry standards.
+
+This package implements the Academy Color Encoding System (ACES), a color management
+framework developed by the Academy of Motion Picture Arts and Sciences. The implementation
+follows the ACES 1.2 standard and is compatible with industry-standard VFX applications
+like Nuke, Houdini, and Maya.
+
+Main Modules
+------------
+transform : High-level IDT/ODT functions
+    rec709_to_aces2065_1 : Rec.709 → ACES2065-1 (Default IDT)
+    aces2065_1_to_rec709 : ACES2065-1 → Rec.709 (Default ODT)
+
+encoding : ACES color space conversions
+    aces2065_1_to_acescct : ACES2065-1 → ACEScct (log-encoded AP1)
+    acescct_to_aces2065_1 : ACEScct → ACES2065-1
+    aces2065_1_to_acescg : ACES2065-1 → ACEScg (linear AP1)
+    acescg_to_aces2065_1 : ACEScg → ACES2065-1
+
+matrix : Low-level matrix transformations
+eotf : Transfer functions (gamma encoding/decoding)
+
+Workflow
+--------
+The default workflow follows the "Utility - sRGB - Texture" IDT approach:
+
+1. Input: Display-referred content (textures, graphics, graded footage)
+2. IDT: rec709_to_aces2065_1() - Gamma decode + matrix transform
+3. Working space: ACES2065-1 (AP0, linear) or ACEScg (AP1, linear)
+4. ODT: aces2065_1_to_rec709() - Matrix transform + gamma encode
+5. Output: Display-ready Rec.709 content
+
+White Level
+-----------
+This implementation maintains normalized white levels (≈1.0) for display-referred
+content, which is the standard for CG/VFX workflows. This differs from some video
+workflows that use elevated white levels (≈16.3) for scene-referred content.
+
+References
+----------
+- ACES Project: https://acescentral.com
+- ACES Technical Documentation: https://docs.acescentral.com
+- OpenColorIO ACES Configs: https://github.com/AcademySoftwareFoundation/OpenColorIO-Config-ACES
+- SMPTE ST 2065-1: ACES2065-1 specification
+"""
+
+from .encoding import (
+    aces2065_1_to_acescct,
+    aces2065_1_to_acescg,
+    acescct_to_aces2065_1,
+    acescg_to_aces2065_1,
+)
+from .transform import aces2065_1_to_rec709, rec709_to_aces2065_1
+
+__all__ = [
+    # Main IDT/ODT
+    "rec709_to_aces2065_1",
+    "aces2065_1_to_rec709",
+    # Color space conversions
+    "aces2065_1_to_acescct",
+    "acescct_to_aces2065_1",
+    "aces2065_1_to_acescg",
+    "acescg_to_aces2065_1",
+]

pixtreme/color/aces/encoding.py

@@ -0,0 +1,186 @@
+"""
+ACES log encoding and color space transformations (ACEScct, ACEScg).
+
+ACEScg (AP1):
+    Linear working space for CG/VFX, using AP1 primaries.
+
+ACEScct (AP1 Log):
+    Log-encoded version of ACEScg, commonly used in color grading.
+    Uses a log curve optimized for grading controls.
+
+References:
+- SMPTE S-2016-001: ACEScg specification
+- Academy TB-2014-004: ACEScct specification
+"""
+
+import cupy as cp
+import numpy as np
+
+from ...utils.dlpack import to_cupy
+from ...utils.dtypes import to_float32
+
+# AP0 (ACES2065-1) ↔ AP1 (ACEScg) transformation matrices
+ap0_to_ap1_matrix = cp.array(
+    [
+        [1.4514393161, -0.2365107469, -0.2149285693],
+        [-0.0765537734, 1.1762296998, -0.0996759264],
+        [0.0083161484, -0.0060324498, 0.9977163014],
+    ]
+)
+
+ap1_to_ap0_matrix = cp.array(
+    [
+        [0.6954522414, 0.1406786965, 0.1638690622],
+        [0.0447945634, 0.8596711185, 0.0955343182],
+        [-0.0055258826, 0.0040252103, 1.0015006723],
+    ]
+)
+
+
+def aces2065_1_to_acescct(image: cp.ndarray | np.ndarray) -> cp.ndarray | np.ndarray:
+    """
+    Convert ACES2065-1 (AP0) to ACEScct (AP1 log-encoded).
+
+    ACEScct uses a log curve optimized for color grading, commonly used
+    in DaVinci Resolve and other grading applications.
+
+    Parameters
+    ----------
+    image : cp.ndarray | np.ndarray
+        ACES2065-1 image (AP0, linear)
+
+    Returns
+    -------
+    cp.ndarray | np.ndarray
+        ACEScct image (AP1, log-encoded)
+    """
+    image = to_float32(image, clip=False)
+
+    # Remember input type
+    was_numpy = isinstance(image, np.ndarray)
+
+    # Convert to CuPy for processing
+    if was_numpy:
+        image = to_cupy(image)
+
+    # AP0 → AP1
+    ap1_image = cp.dot(image, ap0_to_ap1_matrix.T)
+
+    # Apply ACEScct log curve
+    acescct_image = cp.where(
+        ap1_image <= 0.0078125,
+        10.5402377416545 * ap1_image + 0.0729055341958355,
+        (cp.log2(ap1_image) + 9.72) / 17.52,
+    )
+
+    # Convert back to original type
+    if was_numpy:
+        return cp.asnumpy(acescct_image)
+    return acescct_image
+
+
+def acescct_to_aces2065_1(image: cp.ndarray | np.ndarray) -> cp.ndarray | np.ndarray:
+    """
+    Convert ACEScct (AP1 log-encoded) to ACES2065-1 (AP0).
+
+    Parameters
+    ----------
+    image : cp.ndarray | np.ndarray
+        ACEScct image (AP1, log-encoded)
+
+    Returns
+    -------
+    cp.ndarray | np.ndarray
+        ACES2065-1 image (AP0, linear)
+    """
+    image = to_float32(image, clip=False)
+
+    # Remember input type
+    was_numpy = isinstance(image, np.ndarray)
+
+    # Convert to CuPy for processing
+    if was_numpy:
+        image = to_cupy(image)
+
+    # Apply inverse ACEScct log curve
+    ap1_image = cp.where(
+        image <= 0.155251141552511,
+        (image - 0.0729055341958355) / 10.5402377416545,
+        cp.where(
+            image < cp.log2(65504),
+            2 ** ((image * 17.52) - 9.72),
+            65504,
+        ),
+    )
+
+    # AP1 → AP0
+    ap0_image = cp.dot(ap1_image, ap1_to_ap0_matrix.T)
+
+    # Convert back to original type
+    if was_numpy:
+        return cp.asnumpy(ap0_image)
+    return ap0_image
+
+
+def aces2065_1_to_acescg(image: cp.ndarray | np.ndarray) -> cp.ndarray | np.ndarray:
+    """
+    Convert ACES2065-1 (AP0) to ACEScg (AP1 linear).
+
+    ACEScg is the linear working space for CG/VFX using AP1 primaries.
+
+    Parameters
+    ----------
+    image : cp.ndarray | np.ndarray
+        ACES2065-1 image (AP0, linear)
+
+    Returns
+    -------
+    cp.ndarray | np.ndarray
+        ACEScg image (AP1, linear)
+    """
+    image = to_float32(image, clip=False)
+
+    # Remember input type
+    was_numpy = isinstance(image, np.ndarray)
+
+    # Convert to CuPy for processing
+    if was_numpy:
+        image = to_cupy(image)
+
+    ap1_image = cp.dot(image, ap0_to_ap1_matrix.T)
+
+    # Convert back to original type
+    if was_numpy:
+        return cp.asnumpy(ap1_image)
+    return ap1_image
+
+
+def acescg_to_aces2065_1(image: cp.ndarray | np.ndarray) -> cp.ndarray | np.ndarray:
+    """
+    Convert ACEScg (AP1 linear) to ACES2065-1 (AP0).
+
+    Parameters
+    ----------
+    image : cp.ndarray | np.ndarray
+        ACEScg image (AP1, linear)
+
+    Returns
+    -------
+    cp.ndarray | np.ndarray
+        ACES2065-1 image (AP0, linear)
+    """
+    image = to_float32(image, clip=False)
+
+    # Remember input type
+    was_numpy = isinstance(image, np.ndarray)
+
+    # Convert to CuPy for processing
+    if was_numpy:
+        image = to_cupy(image)
+
+    ap0_image = cp.dot(image, ap1_to_ap0_matrix.T)
+
+    # Convert back to original type
+    if was_numpy:
+        return cp.asnumpy(ap0_image)
+    return ap0_image