mvmulti.py

""" Since feisty2's only supports float, this is a drop-in replacement.
    That said, there were some small changes made and it now uses vs.core over vs.get_core()

    All:         can use floats for thsad, etc without integer mvtools throwing errors
                 everything has **args, so now the opt parameter can be changed
    Analyze:     badrange defaults to -24 when search is 3, 6, or 7
                 both 'mvmulti.Analyse' and 'mvmulti.Analyze' work now
    Recalculate: using blksize=2 on integer clips just passes the vectors unprocessed to avoid error messages
    DegrainN:    setting i.e. planes=[0,1,2] will work in place of plane=4
                 for float 32, setting thsadc/limitc overrides array settings: (thsad=[1,2,2], thsadc=1) --> (thsad=[1,1,1])
                 for integers, not setting thsadc/limitc will use the trailing (param[-1]) values from thsad/limit instead
                 trying to use large temporal radii on integer clips will silently fall back to mv.Degrain3
"""

import vapoursynth as vs

def Analyse(super, blksize=8, blksizev=None, levels=0, search=4, searchparam=2, pelsearch=0, _lambda=None, chroma=True, tr=3, truemotion=True, lsad=None, plevel=None, _global=None, pnew=None, pzero=None, pglobal=0, overlap=0, overlapv=None, divide=0, badsad=10000.0, badrange=None, meander=True, trymany=False, fields=False, tff=None, search_coarse=3, dct=0, **args):
    core         = vs.core
    if super.format.sample_type==vs.FLOAT:
        mv       = core.mvsf
    else:
        mv       = core.mv
        badsad   = round(badsad)
    if badrange is None:
        badrange = -24 if search in [3,6,7] else 24
    def getvecs(isb, delta):
        vectors  = mv.Analyse(super, isb=isb, blksize=blksize, blksizev=blksizev, levels=levels, search=search, searchparam=searchparam, pelsearch=pelsearch, _lambda=_lambda, chroma=chroma, delta=delta, truemotion=truemotion, lsad=lsad, plevel=plevel, _global=_global, pnew=pnew, pzero=pzero, pglobal=pglobal, overlap=overlap, overlapv=overlapv, divide=divide, badsad=badsad, badrange=badrange, meander=meander, trymany=trymany, fields=fields, tff=tff, search_coarse=search_coarse, dct=dct, **args)
        return vectors
    bv           = [getvecs(True, i) for i in range(tr, 0, -1)]
    fv           = [getvecs(False, i) for i in range(1, tr+1)]
    vmulti       = bv + fv
    vmulti       = core.std.Interleave(vmulti)
    return vmulti

Analyze = Analyse

def Recalculate(super, vectors, thsad=200.0, smooth=1, blksize=8, blksizev=None, search=4, searchparam=2, _lambda=None, chroma=True, truemotion=True, pnew=None, overlap=0, overlapv=None, divide=0, meander=True, fields=False, tff=None, dct=0, tr=3, **args):
    core         = vs.core
    if super.format.sample_type==vs.FLOAT:
        mv       = core.mvsf
    elif blksize==2:
        return vectors
    else:
        mv       = core.mv
        thsad    = round(thsad)
    def refine(delta):
        analyzed = vectors[delta::2*tr]
        refined  = mv.Recalculate(super, analyzed, thsad=thsad, smooth=smooth, blksize=blksize, blksizev=blksizev, search=search, searchparam=searchparam, _lambda=_lambda, chroma=chroma, truemotion=truemotion, pnew=pnew, overlap=overlap, overlapv=overlapv, divide=divide, meander=meander, fields=fields, tff=tff, dct=dct, **args)
        return refined
    vmulti       = [refine(i) for i in range(0, 2*tr)]
    vmulti       = core.std.Interleave(vmulti)
    return vmulti

def StoreVect(vectors, log):
    core         = vs.core
    w            = vectors.get_frame(0).width
    with open(log, "w") as f:
         print(w, file=f)
    vectors      = core.std.CropAbs(vectors, width=w, height=1)
    return vectors

def RestoreVect(store, log):
    core         = vs.core
    with open(log, "r") as f:
         w       = int(f.readline())
    vectors      = core.raws.Source(store, w, 1, src_fmt="Y8")
    blank        = core.std.BlankClip(vectors, width=1, length=1)
    vectors      = core.std.Splice([blank, vectors], mismatch=True)
    vectors      = core.std.Trim(vectors, 1)
    return vectors

def Compensate(clip, super, vectors, scbehavior=1, thsad=10000.0, fields=False, time=100.0, thscd1=400.0, thscd2=130.0, tff=None, tr=3, cclip=None, **args):
    core         = vs.core
    if super.format.sample_type==vs.FLOAT:
        mv       = core.mvsf
    else:
        mv       = core.mv
        thsad    = round(thsad)
        thscd1   = round(thscd1)
        thscd2   = round(thscd2)
    cclip        = clip if cclip is None else cclip
    def comp(delta):
        v        = vectors[delta::2*tr]
        mc       = mv.Compensate(clip, super, v, scbehavior=scbehavior, thsad=thsad, fields=fields, time=time, thscd1=thscd1, thscd2=thscd2, tff=tff, **args)
        return mc
    bcomp        = [comp(i) for i in range(0, tr)]
    fcomp        = [comp(i) for i in range(tr, 2*tr)]
    compmulti    = bcomp + [cclip] + fcomp
    compmulti    = core.std.Interleave(compmulti)
    return compmulti

def Flow(clip, super, vectors, time=100.0, mode=0, fields=False, thscd1=400.0, thscd2=130.0, tff=None, tr=3, cclip=None, **args):
    core         = vs.core
    if super.format.sample_type==vs.FLOAT:
        mv       = core.mvsf
    else:
        mv       = core.mv
        thscd1   = round(thscd1)
        thscd2   = round(thscd2)
    cclip        = clip if cclip is None else cclip
    def flow(delta):
        v        = vectors[delta::2*tr]
        mc       = mv.Flow(clip, super, v, time=time, mode=mode, fields=fields, thscd1=thscd1, thscd2=thscd2, tff=tff)
        return mc
    bflow        = [flow(i) for i in range(0, tr)]
    fflow        = [flow(i) for i in range(tr, 2*tr)]
    flowmulti    = bflow + [cclip] + fflow
    flowmulti    = core.std.Interleave(flowmulti)
    return flowmulti

def DegrainN(clip, super, mvmulti, tr=3, thsad=400.0, plane=4, limit=None, thscd1=400.0, thscd2=130.0, **args):
    core         = vs.core

    if tr > 24:
        raise ValueError("DegrainN: dude, tr gotta be an int between 1-24, try something less wild maybe?")

    # Some parameter parsing to deal with mvtools using chroma parameters and mvtools-sf using int arrays
    thsad    = thsad if isinstance(thsad, list) else [thsad]
    while len(thsad) < 3:
        thsad += [thsad[-1]]
    limit    = limit if isinstance(limit, list) else [limit] if limit is not None else [1] if clip.format.sample_type==vs.FLOAT else [255]
    while len(limit) < 3:
        limit += [limit[-1]]

    # Just 'cuz
    if 'planes' in args:
        if planes in [0,[0],(0)]:
            plane = 0
        elif planes in [1,[1],(1)]:
            plane = 1
        elif planes in [2,[2],(2)]:
            plane = 2
        elif planes in [[1,2],(1,2)]:
            plane = 3
        else:
            plane = 4

    if clip.format.sample_type==vs.FLOAT:
        mv       = core.mvsf
        tr_ne    = tr
        if 'limitc' in args:
            limit = [limit[0]] + [args['limitc']]*2
        if 'thsadc' in args:
            thsad = [thsad[0]] + [args['thsadc']]*2
    else:
        mv       = core.mv
        tr_ne    = min(3, tr)
        if 'limitc' not in args:
            args.update(thsadc=round(limit[-1]))
        else:
            args['limitc'] = round(args['limitc'])
        if 'thsadc' not in args:
            args.update(thsadc=round(thsad[-1]))
        else:
            args['thsadc'] = round(args['thsadc'])
        if tr > 3:
            mvmulti = core.std.SelectEvery(mvmulti, tr*2, list(range(tr-3,tr*3-3)))
            tr = 3
        thsad    = round(thsad[0])
        limit    = round(limit[0])
        thscd1   = round(thscd1)
        thscd2   = round(thscd2)

    def bvn(n): return mvmulti[tr-n::tr*2]
    def fvn(n): return mvmulti[tr+n-1::tr*2]

    vectors = dict(mvbw=bvn(1), mvfw=fvn(1))
    for i in range(2, tr_ne+1):
        vectors.update({'mvbw{}'.format(i): bvn(i), 'mvfw{}'.format(i): fvn(i)})

    dgn = eval('mv.Degrain{}(clip, super, thsad=thsad, plane=plane, limit=limit, thscd1=thscd1, thscd2=thscd2, **vectors, **args)'.format(tr_ne))

    return dgn

### TODO: make Analyse and Recalculate output arrays like [[b1,b2],[f1,f2]] so you can use mixed tr settings with the same set of vectors