rewrite vine

vsdehalo/__init__.py

@@ -1,6 +1,6 @@
 # ruff: noqa: F401, F403
 
-from . import alpha, denoise, vine, warp
+from . import alpha, denoise, warp
 from .alpha import *
 from .denoise import *
 from .mask import *

vsdehalo/denoise.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 
-from math import ceil
+from functools import partial
+from math import ceil, log
 
 from vsaa import Nnedi3
-from vsdenoise import Prefilter
+from vsdenoise import Prefilter, nl_means, frequency_merge
 from vsexprtools import ExprOp, ExprToken, norm_expr
-from vskernels import NoShift, Point, Scaler, ScalerT
+from vskernels import NoShift, Point, Catrom, Scaler, ScalerT
 from vsmasktools import Morpho, Prewitt
-from vsrgtools import LimitFilterMode, contrasharpening, contrasharpening_dehalo, limit_filter, repair
+from vsrgtools import LimitFilterMode, contrasharpening, contrasharpening_dehalo, limit_filter, repair, gauss_blur
 from vstools import (
-    FieldBased, FunctionUtil, PlanesT, UnsupportedFieldBasedError, check_ref_clip, fallback, mod4, plane, vs
+    FieldBased, FunctionUtil, PlanesT, UnsupportedFieldBasedError, check_ref_clip, fallback, mod4, plane, vs, core
 )
 
 __all__ = [
-    'smooth_dering'
+    'smooth_dering',
+    'vine_dehalo'
 ]
 
 
@@ -153,3 +155,58 @@ def smooth_dering(
         dering = pre_downscaler.scale(work_clip, clip.width, clip.height)
 
     return func.return_clip(dering)
+
+
+def vine_dehalo(
+    clip: vs.VideoNode, strength: int = 16, sharp: float = 0.5, sigma: float = 1.0,
+    supersampler: ScalerT = Nnedi3, downscaler: ScalerT = Catrom, planes: PlanesT = 0, **kwargs
+) -> vs.VideoNode:
+    """
+    :param clip:            Clip to process.
+    :param strength:        Strength of nl_means filtering.
+    :param sharp:           Weight to blend supersampled clip.
+    :param sigma:           Gaussian sigma for filtering cutoff.
+    :param supersampler:    Scaler used for supersampling before dehaloing.
+    :param downscaler:      Scaler used for downscaling after supersampling.
+    :param planes:          Planes to be processed.
+    :param kwargs:          Additional kwargs to be passed to nl_means.
+
+    :return:                Dehaloed clip.
+    """
+    func = FunctionUtil(clip, vine_dehalo, planes)
+
+    if FieldBased.from_video(clip).is_inter:
+        raise UnsupportedFieldBasedError('Only progressive video is supported!', func.func)
+
+    # Only God knows how these were derived.
+    constants = (
+        0.3926327792690057290863679493724,
+        18.880334973195822973214959957208,
+        0.5862453661304626725671053478676
+    )
+
+    sharp = min(max(sharp, 0.0), 1.0)
+    simr = kwargs.pop('simr', None)
+
+    weight = constants[0] * sharp * log(1 + 1 / (constants[0] * sharp))
+    h_refine = constants[1] * (strength / constants[1]) ** constants[2]
+
+    supersampled = supersampler.multi(func.work_clip)
+    supersampled = nl_means(supersampled, strength, tr=0, simr=0, **kwargs)
+    supersampled = downscaler.scale(supersampled, func.work_clip.width, func.work_clip.height)
+
+    smoothed = nl_means(func.work_clip, strength, tr=0, simr=0, **kwargs)
+    smoothed = core.std.Merge(supersampled, smoothed, weight)
+
+    highpassed = frequency_merge(
+        [func.work_clip, smoothed],
+        mode_low=func.work_clip,
+        mode_high=smoothed,
+        lowpass=partial(gauss_blur, sigma=sigma)
+    )
+
+    refine = func.work_clip.std.MakeDiff(highpassed)
+    refine = nl_means(refine, h_refine, tr=0, simr=simr, **kwargs)
+    refine = highpassed.std.MergeDiff(refine)
+
+    return func.return_clip(refine)