From 599b90d2967d3842acad16dfc1b4b0c9c67f92e5 Mon Sep 17 00:00:00 2001 From: emotion3459 <176516814+emotion3459@users.noreply.github.com> Date: Wed, 11 Sep 2024 20:41:53 -0400 Subject: [PATCH] Add VIVTC wrappers (#39) * Add VIVTC wrappers * fix typo * correct vfm return * fix flake8 errors * Update ivtc.py * Add VFMMode enum, vfm docstrings, refactor * vfm: move kwargs merge down * Reduce `vfm_kwargs` dict * Update block/y kwargs, postprocess input clip * Remove field from kwargs * doc * refactor * fixes * Update ivtc.py * Update ivtc.py * Update ivtc.py * cleanup * Docstring, early exit, linting * Remove range in funcutil * pop dryrun --------- Co-authored-by: LightArrowsEXE --- vsdeinterlace/ivtc.py | 155 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 154 insertions(+), 1 deletion(-) diff --git a/vsdeinterlace/ivtc.py b/vsdeinterlace/ivtc.py index 40bb252..db38afb 100644 --- a/vsdeinterlace/ivtc.py +++ b/vsdeinterlace/ivtc.py @@ -2,16 +2,53 @@ from typing import Any -from vstools import CustomEnum, FieldBased, FieldBasedT, InvalidFramerateError, VSFunctionKwArgs, core, join, vs +from vstools import ( + CustomEnum, CustomIntEnum, FieldBased, FieldBasedT, + FunctionUtil, InvalidFramerateError, VSFunctionKwArgs, + VSFunctionNoArgs, core, find_prop_rfs, join, vs +) from .blending import deblend __all__ = [ 'IVTCycles', 'sivtc', 'jivtc', + 'vfm', 'VFMMode', + 'vdecimate' ] +class VFMMode(CustomIntEnum): + """ + Enum representing different matching modes for VFM. + + The mode determines the strategy used for matching fields and frames. + Higher modes generally offer better matching in complex scenarios but + may introduce more risk of jerkiness or duplicate frames. + """ + + TWO_WAY_MATCH = 0 + """2-way match (p/c). Safest option, but may output combed frames in cases of bad edits or blended fields.""" + + TWO_WAY_MATCH_THIRD_COMBED = 1 + """2-way match + 3rd match on combed (p/c + n). Default mode.""" + + TWO_WAY_MATCH_THIRD_SAME_ORDER = 2 + """2-way match + 3rd match (same order) on combed (p/c + u).""" + + TWO_WAY_MATCH_THIRD_FOURTH_FIFTH = 3 + """2-way match + 3rd match on combed + 4th/5th matches if still combed (p/c + n + u/b).""" + + THREE_WAY_MATCH = 4 + """3-way match (p/c/n).""" + + THREE_WAY_MATCH_FOURTH_FIFTH = 5 + """ + 3-way match + 4th/5th matches on combed (p/c/n + u/b). + Highest risk of jerkiness but best at finding good matches. + """ + + class IVTCycles(list[int], CustomEnum): cycle_10 = [[0, 3, 6, 8], [0, 2, 5, 8], [0, 2, 4, 7], [2, 4, 6, 9], [1, 4, 6, 8]] cycle_08 = [[0, 3, 4, 6], [0, 2, 5, 6], [0, 2, 4, 7], [0, 2, 4, 7], [1, 2, 4, 6]] @@ -90,3 +127,119 @@ def jivtc( final = join(ivtced, final) if chroma_only else final return FieldBased.ensure_presence(final, FieldBased.PROGRESSIVE) + + +def vfm( + clip: vs.VideoNode, tff: FieldBasedT | None = None, + mode: VFMMode = VFMMode.TWO_WAY_MATCH_THIRD_COMBED, + postprocess: vs.VideoNode | VSFunctionNoArgs | None = None, + **kwargs: Any +) -> vs.VideoNode: + """ + Perform field matching using VFM. + + This function uses VIVTC's VFM plugin to detect and match pairs of fields in telecined content. + + :param clip: Input clip to field matching telecine on. + :param tff: Field order of the input clip. + If None, it will be automatically detected. + :param mode: VFM matching mode. For more information, see :py:class:`VFMMode`. + Default: VFMMode.TWO_WAY_MATCH_THIRD_COMBED. + :param postprocess: Optional function or clip to process combed frames. + If a function is passed, it should take a clip as input and return a clip as output. + If a clip is passed, it will be used as the postprocessed clip. + :param kwargs: Additional keyword arguments to pass to VFM. + For a list of parameters, see the VIVTC documentation. + + :return: Field matched clip with progressive frames. + """ + + func = FunctionUtil(clip, vfm, None, (vs.YUV, vs.GRAY), 8) + + tff = FieldBased.from_param_or_video(tff, clip, False, func.func) + + vfm_kwargs = dict[str, Any]( + order=tff.is_tff, mode=mode + ) + + if block := kwargs.pop('block', None): + if isinstance(block, int): + vfm_kwargs |= dict(blockx=block, blocky=block) + else: + vfm_kwargs |= dict(blockx=block[0], blocky=block[1]) + + if y := kwargs.pop('y', None): + if isinstance(y, int): + vfm_kwargs |= dict(y0=y, y1=y) + else: + vfm_kwargs |= dict(y0=y[0], y1=y[1]) + + if not kwargs.get('clip2', None) and func.work_clip.format is not clip.format: + vfm_kwargs |= dict(clip2=clip) + + fieldmatch = func.work_clip.vivtc.VFM(**(vfm_kwargs | kwargs)) + + if postprocess: + if callable(postprocess): + postprocess = postprocess(clip) + + fieldmatch = find_prop_rfs(fieldmatch, postprocess, "_Combed", "==", 1) + + return func.return_clip(fieldmatch) + + +def vdecimate(clip: vs.VideoNode, weight: float = 0.0, **kwargs: Any) -> vs.VideoNode: + """ + Perform frame decimation using VDecimate. + + This function uses VIVTC's VDecimate plugin to remove duplicate frames from telecined content. + It's recommended to use the vfm function before running this. + + :param clip: Input clip to decimate. + :param weight: Weight for frame blending. If > 0, blends frames instead of dropping them. + Default: 0.0 (frames are dropped, not blended). + :param kwargs: Additional keyword arguments to pass to VDecimate. + For a list of parameters, see the VIVTC documentation. + + :return: Decimated clip with duplicate frames removed or blended. + """ + + func = FunctionUtil(clip, vdecimate, None, (vs.YUV, vs.GRAY), (8, 16)) + + vdecimate_kwargs = dict[str, Any]() + + if block := kwargs.pop('block', None): + if isinstance(block, int): + vdecimate_kwargs |= dict(blockx=block, blocky=block) + else: + vdecimate_kwargs |= dict(blockx=block[0], blocky=block[1]) + + if not kwargs.get('clip2', None) and func.work_clip.format is not clip.format: + vdecimate_kwargs |= dict(clip2=clip) + + if kwargs.get('dryrun', None): + weight = 0.0 + + if weight: + vdecimate_kwargs |= dict(dryrun=True) + + avg = clip.std.AverageFrames(weights=[0, 1 - weight, weight]) + + if kwargs.get('dryrun', None): + stats = func.work_clip.vivtc.VDecimate(**(vdecimate_kwargs | kwargs)) + + if not weight: + return stats + + vdecimate_kwargs.pop('dryrun', None) + + splice = find_prop_rfs(clip, avg, "VDecimateDrop", "==", 1, stats) + + if kwargs.get('clip2', None): + vdecimate_kwargs |= dict(clip2=splice) + else: + func.work_clip = splice + + decimate = func.work_clip.vivtc.VDecimate(**(vdecimate_kwargs | kwargs)) + + return func.return_clip(decimate)