Clean up pulldown_credits

vsdeinterlace/funcs.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any
 
+from stgpytools import KwargsT
+from vsdenoise import MVTools
 from vstools import DependencyNotFoundError, FieldBased, FieldBasedT, InvalidFramerateError, check_variable, core, vs
 
 __all__ = [
@@ -13,7 +14,8 @@
 def pulldown_credits(
     clip: vs.VideoNode, frame_ref: int, tff: bool | FieldBasedT | None = None,
     interlaced: bool = True, dec: bool | None = None,
-    bob_clip: vs.VideoNode | None = None, qtgmc_args: dict[str, Any] | None = None
+    bob_clip: vs.VideoNode | None = None, qtgmc_args: KwargsT | None = None,
+    mv_args: KwargsT | None = None
 ) -> vs.VideoNode:
     """
     Deinterlacing function for interlaced credits (60i/30p) on top of telecined video (24p).
@@ -41,6 +43,7 @@ def pulldown_credits(
                                     Framerate must be 60000/1001.
     :param qtgmc_args:              Arguments to pass on to QTGMC.
                                     Accepts any parameter except for FPSDivisor and TFF.
+    :param mv_args:                 Arguments to pass on to MVTools, used for motion compensation.
 
     :return:                        IVTC'd/decimated clip with credits pulled down to 24p.
 
@@ -55,20 +58,14 @@ def pulldown_credits(
     except ModuleNotFoundError:
         raise DependencyNotFoundError(pulldown_credits, 'havsfunc')
 
-    try:
-        from vsdenoise import prefilter_to_full_range
-    except ModuleNotFoundError:
-        from havsfunc import DitherLumaRebuild as prefilter_to_full_range  # type: ignore
-        warnings.warn("pulldown_credits: missing dependency `vsdenoise`!", ImportWarning)
-
-    assert check_variable(clip, "pulldown_credits")
+    assert check_variable(clip, pulldown_credits)
 
     InvalidFramerateError.check(pulldown_credits, clip, (30000, 1001))
 
     tff = FieldBased.from_param_or_video(tff, clip, True, pulldown_credits)
     clip = FieldBased.ensure_presence(clip, tff)
 
-    qtgmc_kwargs = dict[str, Any](
+    qtgmc_kwargs = KwargsT(
         SourceMatch=3, Lossless=2, TR0=2, TR1=2, TR2=3, Preset="Placebo"
     ) | (qtgmc_args or {}) | dict(FPSDivisor=1, TFF=tff.field)
 
@@ -79,7 +76,6 @@ def pulldown_credits(
             warnings.warn("pulldown_credits: 'Fieldmatched clip passed to function! "
                           "dec is set to `True`. If you want to disable this, set `dec=False`!'")
 
-    # motion vector and other values
     field_ref = frame_ref * 2
     frame_ref %= 5
     invpos = (5 - field_ref) % 5
@@ -88,8 +84,7 @@ def pulldown_credits(
     pattern = [0, 1, 0, 0, 1][frame_ref]
     direction = [-1, -1, 1, 1, 1][frame_ref]
 
-    blksize = 16 if clip.width > 1024 or clip.height > 576 else 8
-    overlap = blksize // 2
+    mv_args = KwargsT(range_conversion=1.0, pel=2, refine=0) | (mv_args or KwargsT())
 
     ivtc_fps = dict(fpsnum=24000, fpsden=1001)
     ivtc_fps_div = dict(fpsnum=12000, fpsden=1001)
@@ -99,85 +94,85 @@ def pulldown_credits(
 
     InvalidFramerateError.check(pulldown_credits, bobbed, (60000, 1001))
 
-    if interlaced:  # 60i credits. Start of ABBCD
-        if dec:  # Decimate the clip instead of properly IVTC
-            clean = bobbed.std.SelectEvery(5, [4 - invpos])
+    pos = []
+    assumefps = 0
+
+    cycle = 10 // (1 + interlaced)
+
+    # 60i credits. Start of ABBCD
+    if interlaced:
+        if dec:
+            cleanpos = 4
+            pos = [1, 2]
 
             if invpos > 2:
-                jitter = core.std.AssumeFPS(
-                    bobbed[0] * 2 + bobbed.std.SelectEvery(5, [6 - invpos, 7 - invpos]),
-                    **ivtc_fps)  # type:ignore[arg-type]
+                pos = [6, 7]
+                assumefps = 2
             elif invpos > 1:
-                jitter = core.std.AssumeFPS(
-                    bobbed[0] + bobbed.std.SelectEvery(5, [2 - invpos, 6 - invpos]),
-                    **ivtc_fps)  # type:ignore[arg-type]
-            else:
-                jitter = bobbed.std.SelectEvery(5, [1 - invpos, 2 - invpos])
-        else:  # Properly IVTC
+                pos = [2, 6]
+                assumefps = 1
+        else:
+            cleanpos = 1
+            pos = [3, 4]
+
             if invpos > 1:
-                clean = core.std.AssumeFPS(bobbed[0] + bobbed.std.SelectEvery(5, [6 - invpos]),
-                                           **ivtc_fps_div)  # type:ignore[arg-type]
-            else:
-                clean = bobbed.std.SelectEvery(5, [1 - invpos])
+                cleanpos = 6
+                assumefps = 1
 
             if invpos > 3:
-                jitter = core.std.AssumeFPS(bobbed[0] + bobbed.std.SelectEvery(5, [4 - invpos, 8 - invpos]),
-                                            **ivtc_fps)  # type:ignore[arg-type]
-            else:
-                jitter = bobbed.std.SelectEvery(5, [3 - invpos, 4 - invpos])
-
-        jsup_pre = prefilter_to_full_range(jitter, 1.0).mv.Super(pel=2)
-        jsup = jitter.mv.Super(pel=2, levels=1)
-        vect_f = jsup_pre.mv.Analyse(blksize=blksize, isb=False, delta=1, overlap=overlap)
-        vect_b = jsup_pre.mv.Analyse(blksize=blksize, isb=True, delta=1, overlap=overlap)
-        comp = core.mv.FlowInter(jitter, jsup, vect_b, vect_f)
+                pos = [4, 8]
+                assumefps = 1
+
+        clean = bobbed.std.SelectEvery(cycle, [cleanpos - invpos])
+        jitter = bobbed.std.SelectEvery(cycle, [p - invpos for p in pos])
+
+        if assumefps:
+            jitter = core.std.AssumeFPS(bobbed[0] * assumefps + jitter, **(ivtc_fps_div if cleanpos == 6 else ivtc_fps))
+
+        comp = MVTools(jitter, **mv_args).flow_interpolate()
+
         out = core.std.Interleave([comp[::2], clean] if dec else [clean, comp[::2]])
         offs = 3 if dec else 2
+
         return out[invpos // offs:]
-    else:  # 30i credits
-        if pattern == 0:
-            if offset == -1:
-                c1 = core.std.AssumeFPS(bobbed[0] + bobbed.std.SelectEvery(
-                    10, [2 + offset, 7 + offset, 5 + offset, 10 + offset]), **ivtc_fps)  # type:ignore[arg-type]
-            else:
-                c1 = bobbed.std.SelectEvery(10, [offset, 2 + offset, 7 + offset, 5 + offset])
-
-            if offset == 1:
-                c2 = core.std.Interleave([
-                    bobbed.std.SelectEvery(10, [4]),
-                    bobbed.std.SelectEvery(10, [5]),
-                    bobbed[10:].std.SelectEvery(10, [0]),
-                    bobbed.std.SelectEvery(10, [9])
-                ])
-            else:
-                c2 = bobbed.std.SelectEvery(10, [3 + offset, 4 + offset, 9 + offset, 8 + offset])
-        else:
-            if offset == 1:
-                c1 = core.std.Interleave([
-                    bobbed.std.SelectEvery(10, [3]),
-                    bobbed.std.SelectEvery(10, [5]),
-                    bobbed[10:].std.SelectEvery(10, [0]),
-                    bobbed.std.SelectEvery(10, [8])
-                ])
-            else:
-                c1 = bobbed.std.SelectEvery(10, [2 + offset, 4 + offset, 9 + offset, 7 + offset])
-
-            if offset == -1:
-                c2 = core.std.AssumeFPS(bobbed[0] + bobbed.std.SelectEvery(
-                    10, [1 + offset, 6 + offset, 5 + offset, 10 + offset]), **ivtc_fps)  # type:ignore[arg-type]
-            else:
-                c2 = bobbed.std.SelectEvery(10, [offset, 1 + offset, 6 + offset, 5 + offset])
-
-        super1_pre = prefilter_to_full_range(c1, 1.0).mv.Super(pel=2)
-        super1 = c1.mv.Super(pel=2, levels=1)
-        vect_f1 = super1_pre.mv.Analyse(blksize=blksize, isb=False, delta=1, overlap=overlap)
-        vect_b1 = super1_pre.mv.Analyse(blksize=blksize, isb=True, delta=1, overlap=overlap)
-        fix1 = c1.mv.FlowInter(super1, vect_b1, vect_f1, time=50 + direction * 25).std.SelectEvery(4, [0, 2])
-
-        super2_pre = prefilter_to_full_range(c2, 1.0).mv.Super(pel=2)
-        super2 = c2.mv.Super(pel=2, levels=1)
-        vect_f2 = super2_pre.mv.Analyse(blksize=blksize, isb=False, delta=1, overlap=overlap)
-        vect_b2 = super2_pre.mv.Analyse(blksize=blksize, isb=True, delta=1, overlap=overlap)
-        fix2 = c2.mv.FlowInter(super2, vect_b2, vect_f2).std.SelectEvery(4, [0, 2])
-
-        return core.std.Interleave([fix1, fix2] if pattern == 0 else [fix2, fix1])
+
+    def bb(idx: int, cut: bool = False) -> vs.VideoNode:
+        if cut:
+            return bobbed[cycle:].std.SelectEvery(cycle, [idx])
+        return bobbed.std.SelectEvery(cycle, [idx])
+
+    # 30i credits
+    if pattern == 0:
+        c1pos = [0, 2, 7, 5]
+        c2pos = [3, 4, 9, 8]
+
+        if offset == -1:
+            c1pos = [2, 7, 5, 10]
+
+        if offset == 1:
+            c2pos = []
+            c2 = core.std.Interleave([bb(4), bb(5), bb(0, True), bb(9)])
+    else:
+        c1pos = [2, 4, 9, 7]
+        c2pos = [0, 1, 6, 5]
+
+        if offset == 1:
+            c1pos = []
+            c1 = core.std.Interleave([bb(3), bb(5), bb(0, True), bb(8)])
+
+        if offset == -1:
+            c2pos = [1, 6, 5, 10]
+
+    if c1pos:
+        c1 = bobbed.std.SelectEvery(10, [c + offset for c in c1pos])
+
+    if c2pos:
+        c2 = bobbed.std.SelectEvery(10, [c + offset for c in c2pos])
+
+    if offset == -1:
+        c1, c2 = (core.std.AssumeFPS(bobbed[0] + c, **ivtc_fps) for c in (c1, c2))
+
+    fix1 = MVTools(c1, **mv_args).flow_interpolate(time=50 + direction * 25)
+    fix2 = MVTools(c2, **mv_args).flow_interpolate()
+
+    return core.std.Interleave([fix1[::2], fix2[::2]] if pattern == 0 else [fix2[::2], fix1[::2]])