Edge Cleaner rewrite by Vardë

edge_cleaner.py

import vapoursynth as vs
from math import ceil, floor
from rgvs import repair, removegrain
from vsmask.edge import EdgeDetect, PrewittStd
from vsutil import scale_value, split, depth, Dither, Range as CRange

core = vs.core


def edge_cleaner(
    clip: vs.VideoNode, strength: float = 10, rmode: int = 17,
    hot: bool = False, smode: int = 0, edgemask: EdgeDetect = PrewittStd()
) -> vs.VideoNode:
    if clip.format.color_family not in {vs.YUV, vs.GRAY}:
        raise ValueError('edge_cleaner: format not supported')

    bits = clip.format.bits_per_sample
    peak = (1 << bits) - 1 if clip.format.sample_type == vs.INTEGER else 1.0

    clip_y, *chroma = split(clip)

    if smode > 0:
        strength += 4

    main = padding(clip_y, 6, 6, 6, 6)
    # warpsf is way too slow
    main = depth(main, 16, vs.INTEGER, dither_type=Dither.NONE) if clip.format.sample_type == vs.FLOAT else main
    main = main.warp.AWarpSharp2(blur=1, depth=cround(strength / 2)).std.Crop(6, 6, 6, 6)
    main = depth(main, bits, clip.format.sample_type, dither_type=Dither.NONE)

    main = repair(main, clip_y, rmode)

    mask = edgemask.edgemask(clip_y).std.Expr(
        f'x {scale_value(4, 8, bits, CRange.FULL)} < 0 x {scale_value(32, 8, bits, CRange.FULL)} > {peak} x ? ?'
    ).std.InvertMask().std.Convolution([1] * 9)

    final = core.std.MaskedMerge(clip_y, main, mask)

    if hot:
        final = repair(final, clip_y, 2)
    if smode:
        clean = removegrain(clip_y, 17)
        diff = core.std.MakeDiff(clip_y, clean)
        expr = f'x {scale_value(4, 8, bits, CRange.FULL)} < 0 x {scale_value(16, 8, bits, CRange.FULL)} > {peak} x ? ?'
        mask = edgemask.edgemask(
            diff.std.Levels(scale_value(40, 8, bits, CRange.FULL), scale_value(168, 8, bits, CRange.FULL), 0.35)
        )
        mask = removegrain(mask, 7).std.Expr(expr)
        final = core.std.MaskedMerge(final, clip_y, mask)

    return final


def padding(clip: vs.VideoNode, left: int = 0, right: int = 0, top: int = 0, bottom: int = 0) -> vs.VideoNode:
    return clip.resize.Point(
        clip.width + left + right, clip.height + top + bottom,
        src_left=-left, src_top=-top,
        src_width=clip.width + left + right, src_height=clip.height + top + bottom
    )


def cround(x: float) -> int:
    return floor(x + 0.5) if x > 0 else ceil(x - 0.5)


def mod_x(x: int, val: int | float) -> int:
    return max(x * x, cround(val / x) * x)