New AutoColor adjustment filter

[Selur]

Here's a AutoLevelsYUV444 version:

import vapoursynth as vs
import numpy as np
import ctypes
import cv2

core = vs.core

def AutoLevelsYUV444(
    clip: vs.VideoNode,
    method: int = 4,
    clip_limit: float = 1.0,
    gridsize: int = 8,
    strength: float = 0.5
) -> vs.VideoNode:
    fmt = clip.format
    if fmt.color_family != vs.YUV or fmt.subsampling_w or fmt.subsampling_h:
        raise vs.Error("AutoLevelsYUV444: Only YUV444 formats are allowed")

    bits     = fmt.bits_per_sample
    is_float = fmt.sample_type == vs.FLOAT
    peak     = 1.0 if is_float else (1 << bits) - 1

    # Choose NumPy and ctypes types based on bit depth
    if is_float:
        dtype = np.float32
        ctype = ctypes.c_float
    else:
        dtype = np.uint8 if bits <= 8 else np.uint16
        ctype = ctypes.c_uint8 if bits <= 8 else ctypes.c_uint16

    def _process_plane(plane: np.ndarray, mode: str, algo: int) -> np.ndarray:
        # plane: 2D array with values in [0, peak]
        # → 8-bit for OpenCV
        p8 = np.clip((plane / peak) * 255, 0, 255).astype(np.uint8)
        if mode == "clahe":
            clahe = cv2.createCLAHE(clipLimit=clip_limit,
                                    tileGridSize=(gridsize, gridsize))
            out8 = clahe.apply(p8)
        elif mode == "hist":
            out8 = cv2.equalizeHist(p8)
        else:  # scale
            hist = cv2.calcHist([p8],[0],None,[256],[0,256]).ravel()
            cdf  = hist.cumsum(); total = cdf[-1]
            clipv = clip_limit * total / 200.0
            lo    = np.searchsorted(cdf, clipv)
            hi    = np.searchsorted(cdf, total - clipv)
            alpha = 255 / max(hi - lo, 1)
            beta  = -lo * alpha
            out8  = (cv2.convertScaleAbs(p8, alpha=alpha, beta=beta)
                     if algo==0 else np.clip(p8*alpha+beta,0,255).astype(np.uint8))
        # back to original range
        return ((out8.astype(np.float32) / 255) * peak).astype(dtype)

    def read_plane(f: vs.VideoFrame, idx: int) -> np.ndarray:
        """Reads plane idx as (height×width) array via ctypes.from_address."""
        w, h   = clip.width, clip.height
        ptr     = f.get_read_ptr(idx)             # c_void_p
        addr    = ptr if isinstance(ptr, int) else ptr.value
        buf_len = w * h
        # Create ctypes array and convert
        buf_type = ctype * buf_len
        buf      = buf_type.from_address(addr)
        arr      = np.ctypeslib.as_array(buf)    # 1D array
        return arr.reshape((h, w))

    def write_plane(f: vs.VideoFrame, idx: int, data: np.ndarray):
        """Writes data (h×w) back to plane idx."""
        w, h   = clip.width, clip.height
        ptr     = f.get_write_ptr(idx)
        addr    = ptr if isinstance(ptr, int) else ptr.value
        buf_len = w * h
        buf_type = ctype * buf_len
        buf      = buf_type.from_address(addr)
        arr      = np.ctypeslib.as_array(buf).reshape((h, w))
        arr[:, :] = data

    def selector(n: int, f):
        # f can be VideoFrame or [VideoFrame]; ~> we take f[0] if list
        frame = f if isinstance(f, vs.VideoFrame) else f[0]
        # Read Y, U, V
        y = read_plane(frame, 0)
        u = read_plane(frame, 1)
        v = read_plane(frame, 2)

        # Apply method
        if   method == 0:
            y2 = _process_plane(y,    "clahe", 0)
        elif method == 1:
            y2 = _process_plane(y,    "hist",  0)
            u  = _process_plane(u,    "hist",  0)
            v  = _process_plane(v,    "hist",  0)
        elif method == 2:
            y2 = _process_plane(y,    "clahe", 0)
            u  = _process_plane(u,    "clahe", 0)
            v  = _process_plane(v,    "clahe", 0)
        elif method == 3:
            y1 = _process_plane(y,    "clahe", 0)
            y3 = _process_plane(y,    "hist",  0)
            y2 = ((y1 + y3) / 2).astype(dtype)
        elif method == 4:
            y2 = _process_plane(y,    "scale", 0)
        else:  # method == 5
            y2 = _process_plane(y,    "scale", 1)

        # Blend
        y_out = ((1 - strength)*y + strength*y2).astype(dtype)

        # Build new frame
        fout = frame.copy()
        write_plane(fout, 0, y_out)
        write_plane(fout, 1, u)
        write_plane(fout, 2, v)
        return fout

    # Use ModifyFrame (Selector returns VideoFrame)
    return core.std.ModifyFrame(clip=clip, clips=clip, selector=selector)

Which only works on Y.

Cu Selur