eedi3.py (bugfix, eedi3aa, DogaKoboAA)

import vapoursynth as vs
import mvsfunc as mvf
import muvsfunc as muf
import functools
import math


# "eedi3aa" - supersampled eedi3 anti-aliasing
#
# ss  = supersampling
# ssh = supersampling (height)
#      numbers below 10 are read as supersampling factor, else they work as width/height parameters
#      since ssh defaults to ss, setting ss to the width only will use aspect ratio to set the equivalent ssh
#
# kernel_u: specify resizer for supersampling.
#      0 - Point (eedi3_rpow2 is used) **not actually implemented, it just uses taa's eedi3 setting
#      1 - Bilinear - pls don't
#      2 - Bicubic  [2, b(0), c(0.5)]
#      3 - Lanczos  [3, taps(4)]
#      4 - Spline16
#      5 - Spline36 (default)
#      string: any fmtc ['kernel', taps or a1, a2, a3] or ['nnedi3', nsize, nns, qual]
# kernel_d: specify resizer for supersampling.
#      1 - Bilinear - Please note, the AVS wiki was written like 80 years ago. It's a low-complexity resizer, not an AA filter
#      2 - Bicubic  [2, b(-0.5), c(0.25)]
#      3 - Lanczos  [3, taps(ss*4)]
#      4 - Spline16
#      5 - Spline36 (default)
#      string: any fmtc ['kernel', taps or a1, a2, a3], and no, you can't set invks
#
# a, b, g: alpha beta, & gamma eedi3 parameters
# nrad, mdis: eedi3 parameters. defaults to 2, 20 for anything at or below 720p^2, scaling to 3, 30 at 1080p^2 and so on
#
# exmask: provide external aa mask. default uses the taa default mask: tcanny.TCanny(1.2)

def eedi3aa(clip, ss=2., ssh=None, kernel_u=[5], kernel_d=[5], a=0.5, b=0.2, g=None, nrad=None, mdis=None, exmask=None)
    core = vs.core
    # TODO
    # Downscale output
    # Force internal resizers when string inputs correspond w/ the core kernels

    f = clip.format
    st = f.sample_type
    bits = f.bits_per_sample
    isGRAY = f.color_family==vs.GRAY

    if bits!=32 and st==vs.FLOAT:
        raise TypeError('eedi3.eedi3aa: 16 bit float unsupported.')

    iw = clip.width
    ih = clip.height

    ssw = ss
    ssh = ssh if ssh is not None else ss if ss<10 else round( ssw / (iw/ih) )

    ssw = round(iw*ssw) if ssw<10 else ssw
    ssh = round(ih*ssh) if ssh<10 else ssh

    rfac = max(ssw/128, ssh/72)
    mdis = mdis if mdis is not None else nrad*10 if isinstance(nrad, int) else min(20, round(rfac) )
    nrad = nrad if nrad is not None else min(2, rfac//10)

    alpha = min(1, max(0, a) )
    beta = min(1, max(0, b) )
    beta = beta if (alpha+beta)<=1 else 1-alpha
    gamma = g
    nrad = max(0, min(4, nrad) )
    mdis = max(1, min(40, mdis) )

    kernel_u = [kernel_u] if not isinstance(kernel_u, list) else kernel_u
    kernel_d = [kernel_d] if not isinstance(kernel_d, list) else kernel_d

    if kernel_u[0]<=0:
        try:
            import vsTAAmbk as taa
            return taa.TAAmbk(clip, aatype=2, alpha=alpha, beta=beta, gamma=gamma, nrad=nrad, mdis=mdis, mclip=exmask)
        except ModuleNotFoundError:
            import taa
            return taa.TAAmbkX(clip, aatype='eedi3', alpha=alpha, beta=beta, gamma=gamma, nrad=nrad, mdis=mdis, mclip=exmask)

    src = togray(clip)

    if exmask is None:
        tc = y8(src, 1).tcanny.TCanny(sigma=1.2, t_h=8.0, mode=0).std.Maximum()
        mclip = None if if hasattr(core, 'eedi3m') else tc.std.Maximum()
        aamask = tc.std.Convolution([1]*9)
        aamask = togray(aamask, bits, range=1, dmode=1)
    else:
        exmask = togray(exmask)
        mclip = None if if hasattr(core, 'eedi3m') else y8(exmask, 1, 1).std.Maximum()
        aamask = togray(exmask, bits, 1, 1)

    mclipv = None if if hasattr(core, 'eedi3m') else togray(mclip, resize=[ssw,ssh], kernel=2)
    mcliph = None if if hasattr(core, 'eedi3m') else mclipv.std.Transpose()

    if kernel_u[0]=='nnedi3':
        aa = togray(src, min(16, bits))
        aa = aa.nnedi3.nnedi3(1, True, 0, kernel_u[1], kernel_u[2], kernel_u[3])
        aa = aa.std.Transpose()
        aa = aa.nnedi3.nnedi3(1, True, 0, kernel_u[1], kernel_u[2], kernel_u[3])
        aa = togray(aa, 8, resize=[ssh, ssw, 0.5, -0.5])
    else:
        aa = src.std.Transpose()
        if isinstance(kernel_u[0], int):
            aa = togray(aa, 8, resize=[ssh, ssw], kernel=kernel_u)
        else:
            aa = aa.fmtc.resample(ssh, ssw, kernel=kernel_u[0], taps=kernel_u[1], a1=kernel_u[1], a2=kernel_u[2], a3=kernel_u[3], flt=0)
    aa = eedi3(aa, dh=False, alpha=alpha, beta=beta, gamma=gamma, nrad=nrad, mdis=mdis, mclip=mcliph)
    aa = aa.std.Transpose()
    aa = eedi3(aa, dh=False, alpha=alpha, beta=beta, gamma=gamma, nrad=nrad, mdis=mdis, mclip=mclipv)
    if isinstance(kernel_u[0], int):
        aa = togray(aa, bits, resize=[iw, ih], kernel=kernel_d)
    else:
        aa = aa.fmtc.resample(iw, ih, kernel=kernel_d[0], taps=kernel_d[1], a1=kernel_d[1], a2=kernel_d[2], a3=kernel_d[3], flt=0)

    aa = src.std.MaskedMerge(aa, aamask)

    return aa if isGRAY else muf.MergeChroma(aa, clip)

# "DogaKoboAA" - a banquet of questionable image scaling methods all crammed into one
#
# rfactor: parameter of nnedi3_rpow2 (only "2" and "4" are supported)
# nsize: sharpness vs interpolation window. Range is 0 (sharpest) to 3 (strongest AA)
# nns, qual: speed vs quality parameters. range: 0-4 and 1-2. lower 'qual' before lowering 'nns'
# mixed: when False, overlay debicubic/nnedi3/eedi3 clip over input clip.
#        when True, mix pre-eedi3'ed nnedi3_rpow2 clip with a debicubic+spline64 clip and overlay eedi3'ed lines
def DogaKoboAA(clip, w=1280, h=720, rfactor=4, nsize=0, nns=3, qual=2, mixed=False):
    core = vs.core
    # TODO
    # Downscale to resolutions between (iw, ih) and (w, h)
    # Support float resolutions (src_width/height parameter for descale where?)

    f = clip.format
    st = f.sample_type
    bits = f.bits_per_sample
    isGRAY = f.color_family==vs.GRAY

    if bits!=32 and st==vs.FLOAT:
        raise TypeError('eedi3.DogaKoboAA: 16 bit float unsupported.')

    y_src = togray(clip)

    tc = y8(y_src, 1).tcanny.TCanny(sigma=1.2, t_h=8.0, mode=0).std.Maximum()
    aamask = tc.std.Convolution([1]*9)
    aamask = togray(aamask, bits, 1, 1)

    mclip = None if if hasattr(core, 'eedi3m') else tc.std.Maximum()
    mclipv = None if if hasattr(core, 'eedi3m') else mclip.resize.Bicubic(iw*2, ih*2)
    mcliph = None if if hasattr(core, 'eedi3m') else mclipv.std.Transpose()

    dsc = y16(y32(y_src).descale.Debicubic(w, h, 0.2, 0.4))

    rpow = dsc.nnedi3.nnedi3(1, True, 0, nsize, nns, qual).std.Transpose().nnedi3.nnedi3(1, True, 0, nsize, nns, qual)

    if mixed:
        nnclip = togray(rpow, resize=[ih, iw, 0.5, -0.5]).std.Transpose()
        spclip = dsc.fmtc.resample(iw, ih, kernel='spline64', flt=0)
        y_src = mvf.LimitFilter(flt=spclip, src=nnclip, thr=1, elast=1.5)
        y_src = togray(y_src, bits)

    if rfactor==4:
        rpow = rpow.nnedi3.nnedi3(0, True, 0, nsize, nns, qual).std.Transpose().nnedi3.nnedi3(1, True, 0, nsize, nns, qual)
        ssclip = togray(rpow, 8, resize=[iw*2, ih*2, -0.5, -0.5]).std.Transpose()
    else:
        ssclip = togray(rpow, 8, resize=[ih*2, iw*2, 0.5, -0.5], kernel=4)
    aa = eedi3(ssclip, dh=False, alpha=0.25, beta=0.6, gamma=0.15, nrad=3, mdis=30, mclip=mcliph)
    aa = aa.std.Transpose()
    aa = eedi3(aa, dh=False, alpha=0.25, beta=0.6, gamma=0.15, nrad=3, mdis=30, mclip=mclipv)
    lines = togray(aa, bits, resize=[iw, ih])

    merged = y_src.std.MaskedMerge(lines, aamask)

    return merged if isGRAY else muf.MergeChroma(merged, clip)

# eedi3 with mclip bug work-around when dh=False
def eedi3(clip, field=1, dh=False, alpha=None, beta=None, gamma=None, nrad=None, mdis=None, mclip=None)
    core = vs.core

    if not isinstance(clip, vs.VideoNode):
        raise TypeError('eedi3.py: clip must be a video clip!')

    if hasattr(core, 'eedi3m'):
        return clip.eedi3m.EEDI3(field, dh, 0, alpha, beta, gamma, nrad, mdis)

    if not isinstance(mclip, vs.VideoNode) or mclip is not None:
        raise TypeError('eedi3.py: mclip must be a video clip!')
    if field!=0 and field!=1:
        raise TypeError('eedi3.py: only single rate deinterlacing is supported')

    f = clip.format
    range = 1 if f.bits_per_sample==16 and f.sample_type==vs.FLOAT else None

    clip = y8(clip, 3, range)

    if mclip is None:
        return clip.eedi3_092.eedi3(field, True, 0, alpha, beta, gamma, nrad, mdis)

    mclip = y8(mclip, 1, 1)

    iw = clip.width
    ih = clip.height
    mw = mclip.width
    mh = mclip.height

    if (iw, ih)!=(mw, mh)
        if (iw, ih)==(mh, mw):
            mclip = core.std.Transpose(mclip)
        else:
            raise TypeError('eedi3.py: incorrect mclip resolution. Attempt to correct with std.Transpose failed.')

    if dh==False:
        mclip = mclip.std.SeparateFields(tff=False).std.SelectEvery(2, field)
        clip = clip.std.SeparateFields(tff=False).std.SelectEvery(2, field)

    return clip.eedi3_092.eedi3(field, True, 0, alpha, beta, gamma, nrad, mdis, mclip=mclip)




# Utility functions
def togray(clip, bits=None, dmode=3, range=None, int16=True, resize=None, kernel=5):
    core = vs.core

    f = clip.format
    cf = f.color_family
    isGRAY = cf==vs.GRAY
    isYUV = cf==vs.YUV
    isRGB = cf==vs.RGB
    isYCOCG = cf==vs.YCOCG
    in_st = f.sample_type
    in_bits = f.bits_per_sample

    iw = clip.width
    ih = clip.height

    bits = in_bits if bits==None else bits
    st = vs.FLOAT if bits==32 else st if (bits, in_bits)==(16, 16) and not int16 else vs.INTEGER

    f16_to_i16 = (in_bits, bits, in_st, st)==(16, 16, vs.FLOAT, vs.INTEGER)

    probablyfull = isRGB or isYCOCG or f16_to_i16
    range = get_range_int(range, probablyfull)

    scaling = [] if resize is None else [resize] if isinstance(resize, int) else resize
    scaling = scaling+[None,None,None,None,None,None]

    ow = iw if scaling[0] is None else scaling[0]
    oh = ih if scaling[1] is None else scaling[1]
    dscl = (iw * ih) > (ow * oh)

    kernel = [kernel] if isinstance(kernel, int) else kernel
    def_bic = ([-0.5, 0.25] if dscl else [0, 0.5]) if len(kernel)==1 else [(0-kernel[1])/2] if dscl else [(1-kernel[1])/2]
    def_lan = [math.ceil((iw/ow + ih/oh)*2), None] if dscl else [4, None]

    kernel = kernel+def_bic if kernel[0]==2 else kernel+def_lan if kernel[0]==3 else kernel+[None, None]

    Resizer = core.resize.Point if kernel[0]==0 else core.resize.Bilinear if kernel[0]==1 else core.resize.Bicubic if kernel[0]==2 else core.resize.Lanczos if kernel[0]==3 else core.resize.Spline16 if kernel[0]==4 else core.resize.Spline36
    res_params = {'width': scaling[0], 'height': scaling[1], 'src_left': scaling[2], 'src_top': scaling[3], 'src_width': scaling[4], 'src_height': scaling[5], 'filter_param_a': kernel[1], 'filter_param_b': kernel[2]}

    if (in_bits, in_st)==(bits, st) and isGRAY and resize is None:
        pass
    elif (in_bits, in_st)==(bits, st) and isYUV and resize is None:
        return clip.std.ShufflePlanes(0, vs.GRAY)
    else:
        format = core.register_format(vs.GRAY, st, bits, 0, 0)
        matrix = None if isGRAY or isYUV else mvf.GetMatrix(clip, id=True)
        dither_type = dmode if isinstance(dmode, str) else 'ordered' if dmode==0 else 'none' if dmode<3 else 'error_diffusion'
        return Resizer(clip, format=format.id, matrix=matrix, dither_type=dither_type, range_in=range, range=range, **res_params)

y = togray
y8 = functools.partial(togray, bits=8)
y16 = functools.partial(togray, bits=16)
y32 = functools.partial(togray, bits=32)