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`.

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

        >>> 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'

    Args:
        clip: Input clip.
        bitdepth: Desired bitdepth of the output clip.
        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.
        range_in: Input pixel range (defaults to input `clip`'s range).
        range_out: Output pixel range (defaults to input `clip`'s range).
        dither_type: Dithering algorithm. Allows overriding default dithering behavior.
            See [Dither][vstools.DitherType].

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

    Returns:
        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.

    Args:
        frame: Input frame.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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.

    Args:
        clip: Input clip.

    Returns:
        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`.

    Args:
        clip: Input clip.
        insert: Clip to insert into the input clip.
        start_frame: Frame to start inserting from.
        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.

    Returns:
        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:
    """
    Join a list of planes together to form a single clip.

    Args:
        _planes: Planes to combine.

    Returns:
        Combined planes.
    """

    from ..functions import flatten_vnodes, to_arr
    from ..utils import get_color_family, get_video_format

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

    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 |= dict(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,
    planes: PlanesT = None,
    func: FuncExceptT | None = None,
) -> Union[
    ConstantFormatVideoNode,
    Callable[P, ConstantFormatVideoNode],
    Callable[[Callable[P, ConstantFormatVideoNode]], Callable[P, ConstantFormatVideoNode]],
]:
    """
    Wraps `vs-zip <https://github.com/dnjulek/vapoursynth-zip>`.Limiter but only processes
    if clip format is not integer, a min/max val is specified or tv_range is True.

    Args:
        clip_or_func: Clip to process or function that returns a VideoNode to be processed.
        min_val: Lower bound. Defaults to the lowest allowed value for the input. Can be specified for each plane
            individually.
        max_val: Upper bound. Defaults to the highest allowed value for the input. Can be specified for each plane
            individually.
        tv_range: Changes min/max defaults values to LIMITED.
        planes: Planes to process.
        func: Function returned for custom error handling. This should only be set by VS package developers.

    Returns:
        Clamped clip.
    """
    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, planes=planes, 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, planes=planes, 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, planes)

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

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

    Returns:
        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 and clip.format.color_family is vs.RGB:
        clip = vs.core.std.RemoveFrameProps(clip, "_Matrix")

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

def sc_detect(clip: vs.VideoNode, threshold: float = 0.1) -> ConstantFormatVideoNode:
    assert check_variable_format(clip, sc_detect)

    stats = vs.core.std.PlaneStats(shift_clip(clip, -1), clip)

    return vs.core.akarin.PropExpr(
        [clip, stats, stats[1:]],
        lambda: {
            "_SceneChangePrev": f"y.PlaneStatsDiff {threshold} > 1 0 ?",
            "_SceneChangeNext": f"z.PlaneStatsDiff {threshold} > 1 0 ?",
        },
    )

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

    Args:
        clip: Input clip.

    Returns:
        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->...

    Args:
        clips: Sequence of clips to stack recursively.

    Returns:
        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
        ]
    )