decimatch.py

December 13, 2022 ยท View on GitHub

from vapoursynth import core

def bestframeselect(clips, ref, clips2=None, stat_func=core.std.PlaneStats, prop="PlaneStatsDiff", comp_func=min, merge_func=None, debug=False, log=False): """ Picks the 'best' clip(s) for any given frame using stat functions. clips: list of clips for statistics ref: reference clip, e.g. core.average.Mean(clips) / core.median.Median(clips) clips2: list of clips for output, defaults to clips if unset stat_func: function that adds frame properties prop: property added by stat_func to compare comp_func: function to decide which clip to pick, e.g. min, max merge_func: function to merge clips if stat_func returned a list, e.g. core.average.Mean / core.median.Median debug: display values of prop for each clip, and which clip was picked, optionally specify alignment """ from vstools import get_prop

if not clips2:
    clips2 = clips

diffs = [stat_func(clip, ref) for clip in clips]
indices = list(range(len(diffs)))
do_debug, alignment = debug if isinstance(debug, tuple) else (debug, 7)

if log:
    print(f"frame,best,score{',score'.join(map(str, indices))}")

def _select(n, f):
    scores = [
        get_prop(diff.props, prop, float) for diff in f
    ]

    best = comp_func(indices, key=lambda i: scores[i])

    if isinstance(best, list):
        nonlocal merge_func
        if not merge_func:
            if hasattr(ref, "average"):
                merge_func = core.average.Mean
            else: # breaks with >31 clips
                from functools import partial
                merge_func = partial(core.std.AverageFrames, weights=[1] * len(best))
        best_clip = merge_func([clips2[b] for b in best])
        best = "/".join(map(str, best))
    else:
        best_clip = clips2[best]

    if log:
        print(f"{n},{best},{','.join(map(str, scores))}")

    if do_debug:
        return best_clip.text.Text(
            "\n".join([f"Prop: {prop}", *[f"{i}: {s}"for i, s in enumerate(scores)], f"Best: {best}"]), alignment
        )

    return best_clip

return core.std.FrameEval(clips2[0], _select, diffs)

def gen_shifts(clip, n, forward=True, backward=True): shifts = [clip] for cur in range(1, n+1): if forward: shifts.append(clip[cur:]+clip[0]*cur) if backward: shifts.append(clip.std.DuplicateFrames([0]cur)[:-1cur]) return shifts

dvd = core.lsmas.LWLibavSource("dvd.mkv") web = core.lsmas.LWLibavSource("web.mkv")

dvd = bestframeselect(gen_shifts(dvd.std.BoxBlur(), 3), web.std.BoxBlur(), gen_shifts(dvd, 3))

dvd = dvd.std.SelectEvery(cycle=6, offsets=[5]) web = web.std.SelectEvery(cycle=6, offsets=[5])