nlm ¶

@NLMeans
def nl_means(
    clip: vs.VideoNode,
    h: float | Sequence[float] = 1.2,
    tr: int | Sequence[int] = 1,
    a: int | Sequence[int] = 2,
    s: int | Sequence[int] = 4,
    backend: NLMeans.Backend = NLMeans.Backend.AUTO,
    ref: vs.VideoNode | None = None,
    wmode: NLMeans.WeightMode = NLMeans.WeightMode.WELSCH,
    planes: PlanesT = None,
    **kwargs: Any,
) -> ConstantFormatVideoNode:
    """
    Convenience wrapper for NLMeans implementations.

    Filter description [here](https://github.com/Khanattila/KNLMeansCL/wiki/Filter-description).

    Example:
        ```py
        denoised = nl_means(clip, 0.4, backend=nl_means.Backend.CUDA, ...)
        ```

    Args:
        clip: Source clip.
        h: Controls the strength of the filtering. Larger values will remove more noise.
        tr: Temporal Radius. Temporal size = `(2 * d + 1)`. Sets the number of past and future frames to uses for
            denoising the current frame. d=0 uses 1 frame, while d=1 uses 3 frames and so on. Usually, larger values
            result in better denoising. Also known as the `d` parameter.
        a: Search Radius. Spatial size = `(2 * a + 1)^2`. Sets the radius of the search window. a=1 uses 9 pixel, while
            a=2 uses 25 pixels and so on. Usually, larger values result in better denoising.
        s: Similarity Radius. Similarity neighbourhood size = `(2 * s + 1) ** 2`. Sets the radius of the similarity
            neighbourhood window. The impact on performance is low, therefore it depends on the nature of the noise.
        backend: Set the backend to use for processing.
        ref: Reference clip to do weighting calculation. Also known as the `rclip` parameter.
        wmode: Weighting function to use.
        planes: Which planes to process.
        **kwargs: Additional arguments passed to the plugin.

    Returns:
        Denoised clip.
    """

    assert check_variable(clip, nl_means)

    planes = normalize_planes(clip, planes)

    if not planes:
        return clip

    params = dict[str, list[float] | list[int]](h=to_arr(h), d=to_arr(tr), a=to_arr(a), s=to_arr(s))

    # TODO: Remove legacy support for old arguments.
    for sargs, kargs in zip(["strength", "sr", "simr"], ["h", "a", "s"]):
        if sargs in kwargs:
            warnings.warn(f"nl_means: '{sargs}' argument is deprecated, use '{kargs}' instead", DeprecationWarning)
            params[kargs] = to_arr(kwargs.pop(sargs))

    def _nl_means(i: int, channels: str) -> ConstantFormatVideoNode:
        return backend.NLMeans(
            clip,
            **{k: p[i] for k, p in params.items()},
            **{"channels": channels, "rclip": ref, "wmode": wmode, "wref": wmode.wref} | kwargs,
        )

    if clip.format.color_family in {vs.GRAY, vs.RGB}:
        for doc, p in params.items():
            if len(set(p)) > 1:
                warnings.warn(
                    f'nl_means: only "{doc}" first value will be used since clip is {clip.format.color_family.name}',
                    UserWarning,
                )

        return _nl_means(0, "AUTO")

    if (
        all(len(p) < 2 for p in params.values())
        and clip.format.subsampling_w == clip.format.subsampling_h == 0
        and planes == [0, 1, 2]
    ):
        return _nl_means(0, "YUV")

    for k, p in params.items():
        params[k] = normalize_seq(p, 2)

    luma = _nl_means(0, "Y") if 0 in planes else None
    chroma = _nl_means(1, "UV") if 1 in planes or 2 in planes else None

    return join({None: clip, tuple(planes): chroma, 0: luma})