helpers ¶

def pre_ss(  # pyright: ignore[reportInconsistentOverload]
    clip: vs.VideoNode,
    function: VSFunctionNoArgs[vs.VideoNode, vs.VideoNode] | None = None,
    rfactor: float = 2.0,
    sp: type[MixedScalerProcess[_ScalerT, Unpack[_BaseScalerTs]]]
    | MixedScalerProcess[_ScalerT, Unpack[_BaseScalerTs]] = ComplexSuperSamplerProcess[Lanczos],  # type: ignore[assignment]
    supersampler: ScalerLike | Callable[[vs.VideoNode, int, int], vs.VideoNode] | None = None,
    downscaler: ScalerLike | Callable[[vs.VideoNode, int, int], vs.VideoNode] | None = None,
    mod: int = 4,
    planes: Planes = None,
    func: FuncExcept | None = None,
) -> vs.VideoNode:
    """
    Supersamples the input clip, applies a given function to the higher-resolution version,
    and then downscales it back to the original resolution.

    This function generalizes the behavior of
    [SuperSamplerProcess][vsaa.deinterlacers.SuperSamplerProcess] and
    [ComplexSuperSamplerProcess][vsscale.various.ComplexSuperSamplerProcess].

    - Examples:
        ```
        out = pre_ss(clip, lambda clip: cool_function(clip, ...), planes=0)
        ```

        - Passing NNEDI3 as a supersampler:
        ```
        from vsaa import SuperSamplerProcess

        out = pre_ss(clip, lambda clip: cool_function(clip, ...), SuperSamplerProcess)
        ```
        - This works too:
        ```
        from vsaa import SuperSamplerProcess

        out = pre_ss(clip, sp=SuperSamplerProcess(function=lambda clip: cool_function(clip, ...)))
        ```

        - Specifying `supersampler` and `downscaler`:
        ```
        from vskernels import Point

        out = pre_ss(clip, lambda clip: cool_function(clip, ...), supersampler=Point, downscaler=Point, planes=0)
        ```

    Args:
        clip: Source clip.
        function: A function to apply on the supersampled clip.
        rfactor: Scaling factor for supersampling. Defaults to 2.
        sp: A `MixedScalerProcess` instance or class.
            Default is `ComplexSuperSamplerProcess[Lanczos]`.
            It upscales with Lanczos and downscales with Point.
        supersampler: Scaler used to upscale the input clip if `sp` is not specified.
        downscaler: Downscaler used for undoing the upscaling done by the supersampler if `sp` is not specified.
        mod: Ensures the supersampled resolution is a multiple of this value. Defaults to 4.
        planes: Which planes to process.
        func: An optional function to use for error handling.

    Returns:
        A clip with the given function applied at higher resolution, then downscaled back.
    """
    func_util = FunctionUtil(clip, func or pre_ss, planes)

    args = (
        func_util.work_clip,
        mod_x(func_util.work_clip.width * rfactor, mod),
        mod_x(func_util.work_clip.height * rfactor, mod),
    )

    if isinstance(sp, MixedScalerProcess):
        return func_util.return_clip(sp.scale(*args))

    if supersampler and downscaler and function:
        ss = (
            Scaler.ensure_obj(supersampler, func_util.func).scale(*args)
            if is_scaler_like(supersampler)
            else supersampler(*args)
        )

        processed = function(ss)

        args = processed, clip.width, clip.height
        down = (
            Scaler.ensure_obj(downscaler, func_util.func).scale(*args)
            if is_scaler_like(downscaler)
            else downscaler(*args)
        )

        return func_util.return_clip(down)

    if function:
        return pre_ss(clip, sp=sp(function=function), rfactor=rfactor, mod=mod, planes=planes, func=func)

    raise CustomTypeError

def scale_var_clip(
    clip: vs.VideoNode,
    scaler: Scaler | Callable[[Resolution], Scaler] | Callable[[tuple[int, int]], Scaler],
    width: int | Callable[[Resolution], int] | Callable[[tuple[int, int]], int] | None,
    height: int | Callable[[Resolution], int] | Callable[[tuple[int, int]], int],
    shift: tuple[float, float] | Callable[[tuple[int, int]], tuple[float, float]] = (0, 0),
    debug: bool = False,
) -> vs.VideoNode:
    """
    Scale a variable clip to constant or variable resolution.

    Args:
        clip: Source clip.
        scaler: A scaler instance or a callable that returns a scaler instance.
        width: A width integer or a callable that returns the width. If None, it will be calculated from the height and
            the aspect ratio of the clip.
        height: A height integer or a callable that returns the height.
        shift: Optional top shift, left shift tuple or a callable that returns the shifts. Defaults to no shift.
        debug: If True, the `var_width` and `var_height` props will be added to the clip.

    Returns:
        Scaled clip.
    """
    _cached_clips = dict[str, vs.VideoNode]()

    no_accepts_var = list[Scaler]()

    def _eval_scale(f: vs.VideoFrame, n: int) -> vs.VideoNode:
        key = f"{f.width}_{f.height}"

        if key not in _cached_clips:
            res = Resolution(f.width, f.height)

            norm_scaler = scaler(res) if callable(scaler) else scaler
            norm_shift = shift(res) if callable(shift) else shift
            norm_height = height(res) if callable(height) else height

            if width is None:
                norm_width = get_w(norm_height, res.width / res.height)
            else:
                norm_width = width(res) if callable(width) else width

            part_scaler = partial(norm_scaler.scale, width=norm_width, height=norm_height, shift=norm_shift)

            scaled = clip
            if (scaled.width, scaled.height) != (norm_width, norm_height):
                if norm_scaler not in no_accepts_var:
                    try:
                        scaled = part_scaler(clip)
                    except BaseException:
                        no_accepts_var.append(norm_scaler)

                if norm_scaler in no_accepts_var:
                    const_clip = clip.resize.Point(res.width, res.height)

                    scaled = part_scaler(const_clip)

            if debug:
                scaled = scaled.std.SetFrameProps(var_width=res.width, var_height=res.height)

            _cached_clips[key] = scaled

        return _cached_clips[key]

    out_clip = clip if callable(width) or callable(height) else clip.std.BlankClip(width, height)

    return out_clip.std.FrameEval(_eval_scale, clip, clip)