Keyframe generator script

keyframes.py

from argparse import ArgumentParser
from pathlib import Path

from vsexprtools import ExprOp, norm_expr
from vstools import (
    CustomRuntimeError, Keyframes, KwargsT, Sentinel, clip_async_render, core, get_w, mod4, mod_x, ranges_product,
    shift_clip, split, vs
)

name = 'SetsuNoKeyframes'

parser = ArgumentParser(prog=name)
parser.add_argument(
    'video_path', type=Path, help='Video input path.'
)
parser.add_argument(
    '--output', '-o', default='./keyframes.txt', type=Path,
    help='Output path'
)
parser.add_argument(
    '--format', '-f', default=1, type=int,
    help='Whether to use v1 (1) or xvid (-1) format for keyframes output (defaults to v1)'
)
parser.add_argument(
    '--sensitivity', '-s', default=0.225, type=float,
    help=(
        'Sensitivity for deciding if a frame is a scene change. '
        '0 means use default sensitivity with XVID or SCXVID calculation (defaults to 0.225)'
    )
)
parser.add_argument(
    '--blocksize', '-b', default=8, type=int,
    help='Block size for SAD calculation (defaults to 8)'
)
parser.add_argument(
    '--height', '-dh', default=360, type=int,
    help='Height the clip is scaled to. 0 to process at native res. (defaults to 360)'
)
parser.add_argument(
    '--type', '-t', default=3, type=int,
    help=(
        'For sensitivity == 0, select the type of scene change calculation.'
        'WWXD (1), SCXVID (2), Union (3) 1 or 2, Intersection (0) 1 and 2 (defaults to 3)'
    )
)
parser.add_argument(
    '--prefetch', default=0, type=int,
    help=(
        'The prefetch argument defines how many frames should ideally rendered concurrently. '
        '0 defaults to the number of threads. (defaults to 0)'
    )
)
parser.add_argument(
    '--backlog', default=-1, type=int,
    help=(
        'The backlog argument defines how many unconsumed frames the script buffers. '
        '-1 default to the number of threads * 4 (defaults to -1)'
    )
)

core.set_affinity(1.0, core.num_threads << 10, 0)

args = parser.parse_args()

block_size, sensitivity = mod4(args.blocksize), abs(args.sensitivity / 3)

if sensitivity > 0:
    sensitivity -= 5 / 255

video_path = args.video_path.resolve()

if hasattr(core, 'lsmas'):
    clip = core.lsmas.LWLibavSource(video_path)
elif hasattr(core, 'bs'):
    clip = core.bs.VideoSource(video_path)
elif hasattr(core, 'ffms2'):
    clip = core.ffms2.Source(video_path)
else:
    raise CustomRuntimeError(
        'You are missing a source filter! Install one of:\n'
        '\tlsmash - https://github.com/AkarinVS/L-SMASH-Works\n'
        '\tbs     - https://github.com/vapoursynth/bestsource\n'
        '\tffms2  - https://github.com/FFMS/ffms2'
    )

height = mod_x(args.height or clip.height, block_size)
width = get_w(height, clip.width / clip.height, block_size)

as_r_kwargs = KwargsT(prefetch=args.prefetch, backlog=args.backlog)

if sensitivity == 0:
    keyframes = Keyframes.from_clip(clip, args.type, None if args.height <= 0 else width, **as_r_kwargs)
else:
    clip = clip.resize.Bilinear(width, height, vs.YUV444P8)

    diff = norm_expr(
        [split(clip), split(shift_clip(clip, -1))], 'x a - abs y b - abs z c - abs + +', None, vs.GRAY10
    )

    sad = norm_expr(
        diff,
        [
            'X {block_size} % 0 = Y {block_size} % 0 = and ',
            [
                'x' if x == y == 0 else ExprOp.REL_PIX('x', x, y)
                for (x, y) in ranges_product(block_size, block_size)
            ],
            ExprOp.ADD * ((block_size * block_size) - 1),
            block_size * block_size, ExprOp.DIV, 0, ExprOp.TERN
        ], None, block_size=block_size, force_akarin=name
    )

    down = norm_expr(sad, 'X {block_size} * Y {block_size} * x[]', block_size=block_size, force_akarin=name)
    down = down.std.CropAbs(width // block_size, height // block_size)

    blank = core.std.BlankClip(None, 1, 1, vs.GRAY8, down.num_frames, keep=True)

    stats = blank.std.CopyFrameProps(down.std.PlaneStats())

    stats = norm_expr(stats, f'x.PlaneStatsAverage {sensitivity} >', force_akarin=name)

    frames = clip_async_render(
        stats, None, 'Detecting scene changes...', lambda n, f: Sentinel.check(n, f[0][0, 0]),  # type: ignore
        **as_r_kwargs
    )

    keyframes = Keyframes(Sentinel.filter(frames), clip.num_frames)

keyframes.to_file(args.output, args.format)