utils ¶

def depth(
    clip: vs.VideoNode, bitdepth: VideoFormatT | HoldsVideoFormatT | int | None = None, /,
    sample_type: int | vs.SampleType | None = None, *,
    range_in: ColorRangeT | None = None, range_out: ColorRangeT | None = None,
    dither_type: str | DitherType = DitherType.AUTO,
) -> ConstantFormatVideoNode:
    """
    A convenience bitdepth conversion function using only internal plugins if possible.

    This uses exclusively internal plugins except for specific dither_types.
    To check whether your DitherType uses fmtc, use `DitherType.is_fmtc`.

    .. code-block:: python

        >>> src_8 = vs.core.std.BlankClip(format=vs.YUV420P8)
        >>> src_10 = depth(src_8, 10)
        >>> src_10.format.name
        'YUV420P10'

    .. code-block:: python

        >>> src2_10 = vs.core.std.BlankClip(format=vs.RGB30)
        >>> src2_8 = depth(src2_10, 8, dither_type=Dither.RANDOM)  # override default dither behavior
        >>> src2_8.format.name
        'RGB24'

    :param clip:            Input clip.
    :param bitdepth:        Desired bitdepth of the output clip.
    :param sample_type:     Desired sample type of output clip. Allows overriding default float/integer behavior.
                            Accepts ``vapoursynth.SampleType`` enums ``vapoursynth.INTEGER`` and ``vapoursynth.FLOAT``
                            or their values, ``0`` and ``1`` respectively.
    :param range_in:        Input pixel range (defaults to input `clip`'s range).
    :param range_out:       Output pixel range (defaults to input `clip`'s range).
    :param dither_type:     Dithering algorithm. Allows overriding default dithering behavior. See :py:class:`Dither`.

                            When integer output is desired but the conversion may produce fractional values,
                            defaults to :attr:`Dither.VOID` if it is available via the fmtc VapourSynth plugin,
                            or to Floyd-Steinberg :attr:`Dither.ERROR_DIFFUSION` for 8-bit output
                            or :attr:`DitherType.ORDERED` for higher bit depths.
                            In other cases, defaults to :attr:`Dither.NONE`, or round to nearest.
                            See :py:func:`Dither.should_dither` for more information.

    :return:                Converted clip with desired bit depth and sample type.
                            ``ColorFamily`` will be same as input.
    """

    from ..utils import get_video_format

    assert check_variable_format(clip, depth)

    in_fmt = get_video_format(clip)
    out_fmt = get_video_format(bitdepth or clip, sample_type=sample_type)

    range_out = ColorRange.from_param_or_video(range_out, clip)
    range_in = ColorRange.from_param_or_video(range_in, clip)

    if (
        in_fmt.bits_per_sample, in_fmt.sample_type, range_in
    ) == (
        out_fmt.bits_per_sample, out_fmt.sample_type, range_out
    ):
        return clip

    dither_type = DitherType(dither_type)

    if dither_type is DitherType.AUTO:
        should_dither = DitherType.should_dither(in_fmt, out_fmt, range_in, range_out)

        if hasattr(clip, "fmtc"):
            dither_type = DitherType.VOID
        else:
            dither_type = DitherType.ERROR_DIFFUSION if out_fmt.bits_per_sample == 8 else DitherType.ORDERED
        dither_type = dither_type if should_dither else DitherType.NONE

    new_format = in_fmt.replace(
        bits_per_sample=out_fmt.bits_per_sample, sample_type=out_fmt.sample_type
    )

    return dither_type.apply(clip, new_format, range_in, range_out)

def frame2clip(frame: vs.VideoFrame) -> ConstantFormatVideoNode:
    """
    Convert a VideoFrame to a VideoNode.

    :param frame:       Input frame.

    :return:            1-frame long VideoNode of the input frame.
    """

    key = hash((frame.width, frame.height, frame.format.id))

    if _f2c_cache.get(key, None) is None:
        _f2c_cache[key] = blank_clip = vs.core.std.BlankClip(
            None, frame.width, frame.height,
            frame.format.id, 1, 1, 1,
            [0] * frame.format.num_planes,
            True
        )
    else:
        blank_clip = _f2c_cache[key]

    frame_cp = frame.copy()

    return vs.core.std.ModifyFrame(blank_clip, blank_clip, lambda n, f: frame_cp)

def get_b(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the blue plane of the given clip.

    :param clip:                Input clip.

    :return:                    B plane of the input clip.

    :raises CustomValueError:   Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_b)

    return plane(clip, 2)

def get_g(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the green plane of the given clip.

    :param clip:                Input clip.

    :return:                    G plane of the input clip.

    :raises CustomValueError:   Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_g)

    return plane(clip, 1)

def get_r(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the red plane of the given clip.

    :param clip:                Input clip.

    :return:                    R plane of the input clip.

    :raises CustomValueError:   Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_r)

    return plane(clip, 0)

def get_u(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the first chroma (U) plane of the given clip.

    :param clip:                Input clip.

    :return:                    Y plane of the input clip.

    :raises CustomValueError:   Clip is not YUV.
    """

    InvalidColorFamilyError.check(clip, vs.YUV, get_u)

    return plane(clip, 1)

def get_v(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the second chroma (V) plane of the given clip.

    :param clip:                Input clip.

    :return:                    V plane of the input clip.

    :raises CustomValueError:   Clip is not YUV.
    """

    InvalidColorFamilyError.check(clip, vs.YUV, get_v)

    return plane(clip, 2)

def get_y(clip: vs.VideoNode, /) -> ConstantFormatVideoNode:
    """
    Extract the luma (Y) plane of the given clip.

    :param clip:                Input clip.

    :return:                    Y plane of the input clip.

    :raises CustomValueError:   Clip is not GRAY or YUV.
    """

    InvalidColorFamilyError.check(clip, [vs.YUV, vs.GRAY], get_y)

    return plane(clip, 0)

def insert_clip(
    clip: vs.VideoNode, /, insert: vs.VideoNode, start_frame: int, strict: bool = True
) -> ConstantFormatVideoNode:
    """
    Replace frames of a longer clip with those of a shorter one.

    The insert clip may NOT exceed the final frame of the input clip.
    This limitation can be circumvented by setting `strict=False`.

    :param clip:                Input clip.
    :param insert:              Clip to insert into the input clip.
    :param start_frame:         Frame to start inserting from.
    :param strict:              Throw an error if the inserted clip exceeds the final frame of the input clip.
                                If False, truncate the inserted clip instead.
                                Default: True.

    :return:                    Clip with frames replaced by the insert clip.

    :raises CustomValueError:   Insert clip is too long, ``strict=False``,
                                and exceeds the final frame of the input clip.
    """

    if start_frame == 0:
        return vs.core.std.Splice([insert, clip[insert.num_frames:]])

    pre = clip[:start_frame]
    insert_diff = (start_frame + insert.num_frames) - clip.num_frames

    if insert_diff == 0:
        return vs.core.std.Splice([pre, insert])

    if insert_diff < 0:
        return vs.core.std.Splice([pre, insert, clip[insert_diff:]])

    if strict:
        raise ClipLengthError(
            'Inserted clip is too long and exceeds the final frame of the input clip.',
            insert_clip, {'clip': clip.num_frames, 'diff': insert_diff}
        )

    return vs.core.std.Splice([pre, insert[:-insert_diff]])

def join(*_planes: Any, **kwargs: Any) -> vs.VideoNode:
    from ..functions import flatten_vnodes, to_arr
    from ..utils import get_color_family, get_video_format

    family: vs.ColorFamily | None = kwargs.get('family', None)

    if isinstance(_planes[-1], vs.ColorFamily):
        family = _planes[-1]
        _planes = _planes[:-1]

    if isinstance(_planes[0], Mapping):
        planes_map = dict[int, vs.VideoNode]()

        for p_key, node in _planes[0].items():
            if node is None:
                continue

            if p_key is None:
                planes_map |= {i: n for i, n in enumerate(flatten_vnodes(node, split_planes=True))}
            else:
                planes_map |= {i: plane(node, min(i, node.format.num_planes - 1)) for i in to_arr(p_key)}

        return join(*(planes_map[i] for i in sorted(planes_map.keys())))

    planes = list[vs.VideoNode](_planes[0] if (
        isinstance(_planes[0], Iterable) and not isinstance(_planes[0], vs.VideoNode)
    ) else _planes)

    n_clips = len(planes)

    if not n_clips:
        raise CustomIndexError('Not enough clips/planes passed!', join)

    if n_clips == 1 and (family is None or family is (planes[0].format and planes[0].format.color_family)):
        return planes[0]

    if family is None:
        family = get_color_family(planes[0])

    if n_clips == 2:
        other_family = get_color_family(planes[1])

        if family in {vs.GRAY, vs.YUV}:
            InvalidColorFamilyError.check(
                other_family, vs.YUV, join, '"chroma" clip must be {correct} color family, not {wrong}!'
            )

            if family is vs.GRAY:
                planes[0] = get_y(planes[0])

            return vs.core.std.ShufflePlanes(planes, [0, 1, 2], other_family)

    if n_clips in {3, 4}:
        if family is vs.GRAY:
            for pplane in planes[:3]:
                if (fmt := get_video_format(pplane)).num_planes > 1:
                    family = fmt.color_family
                    break
            else:
                matrix = Matrix.from_video(planes[0])
                family = vs.RGB if matrix is Matrix.RGB else vs.YUV

        clip = vs.core.std.ShufflePlanes(planes[:3], [0, 0, 0], family)

        if n_clips == 4:
            clip = clip.std.ClipToProp(planes[-1], '_Alpha')

        return clip
    elif n_clips > 4:
        raise CustomValueError('Too many clips or planes passed!', join)

    raise CustomValueError('Not enough clips or planes passed!', join)

def limiter(
    clip_or_func: vs.VideoNode | Callable[P, ConstantFormatVideoNode] | None = None,
    /,
    min_val: float | Sequence[float] | None = None,
    max_val: float | Sequence[float] | None = None,
    *,
    tv_range: bool = False,
    func: FuncExceptT | None = None
) -> Union[
    ConstantFormatVideoNode,
    Callable[P, ConstantFormatVideoNode],
    Callable[[Callable[P, ConstantFormatVideoNode]], Callable[P, ConstantFormatVideoNode]]
]:
    if callable(clip_or_func):
        _func = clip_or_func

        @wraps(_func)
        def _wrapper(*args: P.args, **kwargs: P.kwargs) -> ConstantFormatVideoNode:
            return limiter(_func(*args, **kwargs), min_val, max_val, tv_range=tv_range, func=func or _func)

        return _wrapper

    func = func or limiter
    clip = clip_or_func

    if clip is None:
        return partial(limiter, min_val=min_val, max_val=max_val, tv_range=tv_range, func=func)

    assert check_variable_format(clip, func)

    if all([
        clip.format.sample_type == vs.INTEGER,
        min_val is None,
        max_val is None,
        tv_range is False
    ]):
        return clip

    if not (min_val == max_val is None):
        from ..utils import get_lowest_values, get_peak_values

        min_val = normalize_seq(min_val or get_lowest_values(clip, clip), clip.format.num_planes)
        max_val = normalize_seq(max_val or get_peak_values(clip, clip), clip.format.num_planes)

    return clip.vszip.Limiter(min_val, max_val, tv_range)

def plane(clip: vs.VideoNode, index: int, /, strict: bool = True) -> ConstantFormatVideoNode:
    """
    Extract a plane from the given clip.

    :param clip:        Input clip.
    :param index:       Index of the plane to extract.

    :return:            Grayscale clip of the clip's plane.
    """

    assert check_variable_format(clip, plane)

    if clip.format.num_planes == 1 and index == 0:
        return clip

    if not strict:
        if clip.format.color_family is vs.RGB:
            clip = vs.core.std.RemoveFrameProps(clip, '_Matrix')

    return vs.core.std.ShufflePlanes(clip, index, vs.GRAY)

def split(clip: vs.VideoNode, /) -> list[ConstantFormatVideoNode]:
    """
    Split a clip into a list of individual planes.

    :param clip:    Input clip.

    :return:        List of individual planes.
    """

    assert check_variable_format(clip, split)

    return [clip] if clip.format.num_planes == 1 else [
        plane(clip, i, False) for i in range(clip.format.num_planes)
    ]

def stack_clips(
    clips: Sequence[vs.VideoNode | Sequence[vs.VideoNode | Sequence[
        vs.VideoNode | Sequence[vs.VideoNode | Sequence[vs.VideoNode | Sequence[vs.VideoNode]]]
    ]]]
) -> vs.VideoNode:
    """
    Stack clips in the following repeating order: hor->ver->hor->ver->...

    :param clips:   Sequence of clips to stack recursively.

    :return:        Stacked clips.
    """

    return vs.core.std.StackHorizontal([
        inner_clips if isinstance(inner_clips, vs.VideoNode) else (
            vs.core.std.StackVertical([
                clipa if isinstance(clipa, vs.VideoNode) else stack_clips(clipa)
                for clipa in inner_clips
            ])
        )
        for inner_clips in clips
    ])