mv and mvsf compatibility layer

mvmulti.py

import itertools
import vapoursynth as vs

def super(clip, *args, **kwargs):
    core = vs.get_core()
    sample_type = clip.format.sample_type
    mv = core.mv if sample_type == 0 else core.mvsf
    
    return mv.Super(clip, *args, **kwargs)
        
def analyse(super, *args, **kwargs):
    core = vs.get_core()
    sample_type = super.format.sample_type
    mv = core.mv if sample_type == 0 else core.mvsf
    
    tr = kwargs.pop('tr', 3)

    assert 'isb' not in kwargs
    assert 'delta' not in kwargs    
    assert tr > 0, "temporal radius (tr) must be greater than 0"
    assert tr <= 6 or (sample_type == 1 and tr <= 24), "temporal radius (tr) must be 1..6 for 8-16 bit clips or 1..24 for 32 bit clips"
    
    def _getvecs(isb, delta):
        vectors = mv.Analyse(super, *args, isb=isb, delta=delta, **kwargs)
        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
    return core.std.Interleave(vmulti)
    
def recalculate(super, vectors, *args, **kwargs):
    core = vs.get_core()
    sample_type = super.format.sample_type
    mv = core.mv if sample_type == 0 else core.mvsf
    
    tr = vectors.num_frames // (super.num_frames * 2)
    assert tr <= 6 or (sample_type == 1 and tr <= 24), "temporal radius (tr) must be 1..6 for 8-16 bit clips or 1..24 for 32 bit clips"
    
    def _refine(delta):
        analyzed = core.std.SelectEvery(vectors, 2 * tr, delta)
        refined = mv.Recalculate(super, analyzed, *args, **kwargs)
        return refined
    
    vmulti = [refine(i) for i in range(0, 2 * tr)]
    return core.std.Interleave(vmulti)
    
def compensate(clip, super, vectors, *args, **kwargs):
    core = vs.get_core()
    sample_type = clip.format.sample_type
    mv = core.mv if sample_type == 0 else core.mvsf
    
    tr = vectors.num_frames // (super.num_frames * 2)
    assert tr <= 6 or (sample_type == 1 and tr <= 24), "temporal radius (tr) must be 1..6 for 8-16 bit clips or 1..24 for 32 bit clips"
    
    def _compensate(delta):
        vec = core.std.SelectEvery(vectors, 2 * tr, delta)
        return mv.Compensate(clip, super, vec, *args, **kwargs)
    
    bcomp = [comp(i) for i in range(0, tr)]
    fcomp = [comp(i) for i in range(tr, 2 * tr)]
    compmulti = bcomp + [clip] + fcomp
    return core.std.Interleave(compmulti)

def flow(clip, super, vectors, *args, **kwargs):
    core = vs.get_core()
    sample_type = clip.format.sample_type
    mv = core.mv if sample_type == 0 else core.mvsf
    
    tr = vectors.num_frames // (super.num_frames * 2)
    assert tr <= 6 or (sample_type == 1 and tr <= 24), "temporal radius (tr) must be 1..6 for 8-16 bit clips or 1..24 for 32 bit clips"
    
    def _flow(delta):
        vec = core.std.SelectEvery(vectors, 2 * tr, delta)
        return mv.Flow(clip, super, vec, *args, **kwargs)
        
    bflow = [_flow(i) for i in range(0, tr)]
    fflow = [_flow(i) for i in range(tr, 2 * tr)]
    flowmulti = bflow + [clip] + fflow
    return core.std.Interleave(flowmulti)

def _degrain_n(clip, super, *vectors, **kwargs):
    core = vs.get_core()
    
    tr = len(vectors) // 2
    
    assert 0 < tr <= 6, "temporal radius (tr) must be 1..6"
    
    v = vectors
    if tr == 1:
        sm = core.mv.Degrain1(clip, super, *v, **kwargs)
    elif tr == 2:
        sm = core.mv.Degrain2(clip, super, *v, **kwargs)
    elif tr == 3:
        sm = core.mv.Degrain3(clip, super, *v, **kwargs)
    elif tr == 4:
        mv12 = core.mv.Degrain2(clip, super, *v[:4], **kwargs)
        mv34 = core.mv.Degrain2(clip, super, *v[4:], **kwargs)
        sm = core.std.Merge(mv12, mv34, 0.4444)
    elif tr == 5:
        mv123 = core.mv.Degrain3(clip, super, *v[:6], **kwargs)
        mv45 = core.mv.Degrain2(clip, super, *v[6:], **kwargs)
        sm = core.std.Merge(mv123, mv45, 0.4545)
    elif tr == 6:
        mv123 = core.mv.Degrain3(clip, super, *v[:6], **kwargs)
        mv456 = core.mv.Degrain3(clip, super, *v[6:], **kwargs)
        sm = core.std.Merge(mv123, mv456, 0.4615)
        
    return sm

def degrain(clip, super, mvmulti, *args, **kwargs):
    core = vs.get_core()
    sample_type = clip.format.sample_type
    
    tr = mvmulti.num_frames // (super.num_frames * 2)
    assert tr <= 6 or (sample_type == 1 and tr <= 24), "temporal radius (tr) must be 1..6 for 8-16 bit clips or 1..24 for 32 bit clips"

    def bvn(n):
        return core.std.SelectEvery(mvmulti, tr * 2, tr - n)
        
    def fvn(n):
        return core.std.SelectEvery(mvmulti, tr * 2, tr + n - 1)
    
    if sample_type == 0:
        # mvsf changed limit and limitc type to float [0.0..1.0]
        # float type is expected for 8..16 bit clips too
        kwargs = kwargs.copy()
        kwargs['limit'] = int(round(kwargs.pop('limit', 1.0) * 255))
        kwargs['limitc'] = int(round(kwargs.pop('limitc', 1.0) * 255))
        _degrain = _degrain_n
    else:
        _degrain = getattr(core.mvsf, 'Degrain{n}'.format(n=tr))

    _zipped_vectors = [(bvn(n), fvn(n)) for n in range(1, tr + 1)]
    _vectors = list(itertools.chain.from_iterable(_zipped_vectors))
    
    return _degrain(clip, super, *_vectors, *args, **kwargs)