From 4ec7b184b9818f62ab97ed577988ebd3c1d83b47 Mon Sep 17 00:00:00 2001
From: louis <louis@poweris.moe>
Date: Sun, 2 May 2021 18:12:42 -0400
Subject: [PATCH] friendship ended with knlmeanscl

it keeps leaking my memory and bm3d works after a kernel upgrade
---
 .../tanteidan_common/filter.py                |  7 +-
 Vivy/vivy_common/filter.py                    | 19 ++++--
 yt_common/yt_common/__init__.py               |  2 +-
 yt_common/yt_common/denoise.py                | 65 +++++++++++++++++++
 4 files changed, 83 insertions(+), 10 deletions(-)
 create mode 100644 yt_common/yt_common/denoise.py

diff --git a/Bishounen Tanteidan/tanteidan_common/filter.py b/Bishounen Tanteidan/tanteidan_common/filter.py
index 71b6fde..2454cae 100644
--- a/Bishounen Tanteidan/tanteidan_common/filter.py	
+++ b/Bishounen Tanteidan/tanteidan_common/filter.py	
@@ -12,10 +12,11 @@ from lvsfunc.misc import replace_ranges
 from lvsfunc.types import Range
 from vardefunc import dumb3kdb
 
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from yt_common.antialiasing import combine_mask, sraa_clamp
 from yt_common.deband import morpho_mask
+from yt_common.denoise import bm3d
 
 from .scenefilter import get_op_scenefilters
 
@@ -23,8 +24,8 @@ from .scenefilter import get_op_scenefilters
 core = vs.core
 
 
-def denoise(clip: vs.VideoNode, h: float = 0.4) -> vs.VideoNode:
-    return clip.knlm.KNLMeansCL(d=3, a=1, h=h)
+def denoise(clip: vs.VideoNode, sigma: Union[float, List[float]] = 0.75) -> vs.VideoNode:
+    return bm3d(clip, sigma=sigma)
 
 
 def deband(clip: vs.VideoNode) -> vs.VideoNode:
diff --git a/Vivy/vivy_common/filter.py b/Vivy/vivy_common/filter.py
index fe9f23a..8b3ed32 100644
--- a/Vivy/vivy_common/filter.py
+++ b/Vivy/vivy_common/filter.py
@@ -5,10 +5,14 @@ import vardefunc as vdf
 
 from awsmfunc import bbmod
 from debandshit import f3kbilateral
+from lvsfunc.aa import upscaled_sraa
+from lvsfunc.kernels import Bicubic
+from lvsfunc.misc import replace_ranges
 from lvsfunc.types import Range
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from yt_common import antialiasing
+from yt_common.denoise import bm3d
 from yt_common.deband import morpho_mask
 
 import os
@@ -39,8 +43,8 @@ def letterbox_edgefix(clip: vs.VideoNode, ranges: List[Range]) -> vs.VideoNode:
     return lvf.misc.replace_ranges(clip, edgefix, ranges)
 
 
-def denoise(clip: vs.VideoNode, h: float = 0.4) -> vs.VideoNode:
-    return clip.knlm.KNLMeansCL(d=3, a=1, h=h)
+def denoise(clip: vs.VideoNode, sigma: Union[float, List[float]] = 0.75) -> vs.VideoNode:
+    return bm3d(clip, sigma=sigma)
 
 
 def deband(clip: vs.VideoNode) -> vs.VideoNode:
@@ -58,9 +62,12 @@ def deband(clip: vs.VideoNode) -> vs.VideoNode:
     return deband
 
 
-def antialias(clip: vs.VideoNode, noaa: Optional[List[Range]] = None) -> vs.VideoNode:
-    clamp = antialiasing.sraa_clamp(clip, mask=antialiasing.mask_strong(clip))
-    return lvf.misc.replace_ranges(clamp, clip, noaa) if noaa else clamp
+def antialias(clip: vs.VideoNode, weak: Optional[List[Range]] = None, strong: Optional[List[Range]] = None,
+              noaa: Optional[List[Range]] = None) -> vs.VideoNode:
+    mask = antialiasing.combine_mask(clip, weak or [])
+    clamp = antialiasing.sraa_clamp(clip, mask=mask)
+    sraa = core.std.MaskedMerge(clip, upscaled_sraa(clip, rfactor=2, downscaler=Bicubic(b=0, c=1/2).scale), mask)
+    return replace_ranges(replace_ranges(clamp, clip, noaa or []), sraa, strong or [])
 
 
 def regrain(clip: vs.VideoNode) -> vs.VideoNode:
diff --git a/yt_common/yt_common/__init__.py b/yt_common/yt_common/__init__.py
index 3007da5..a6a8178 100644
--- a/yt_common/yt_common/__init__.py
+++ b/yt_common/yt_common/__init__.py
@@ -1 +1 @@
-from . import antialiasing, automation, config, deband, logging, source  # noqa: F401
+from . import antialiasing, automation, config, deband, denoise, logging, source  # noqa: F401
diff --git a/yt_common/yt_common/denoise.py b/yt_common/yt_common/denoise.py
new file mode 100644
index 0000000..88ce4a5
--- /dev/null
+++ b/yt_common/yt_common/denoise.py
@@ -0,0 +1,65 @@
+import vapoursynth as vs
+from vsutil import iterate, Range, get_y
+from typing import Any, Dict, List, Optional, Union
+
+
+core = vs.core
+
+
+# always uses float OPP
+# mvsfunc is an unreadable mess and i was having bm3d issues so i rewrote the wrapper
+# turns out, it's still an unreadable mess
+def bm3d(clip: vs.VideoNode, matrix_s: str = "709", sigma: Union[float, List[float]] = 0.75,
+         radius: Union[int, List[int], None] = None, ref: Optional[vs.VideoNode] = None,
+         pre: Optional[vs.VideoNode] = None, refine: int = 1,
+         basic_args: Dict[str, Any] = {}, final_args: Dict[str, Any] = {}) -> vs.VideoNode:
+    assert clip.format is not None
+    isGray = clip.format.color_family == vs.GRAY
+
+    def to_opp(clip: vs.VideoNode) -> vs.VideoNode:
+        return clip.resize.Bicubic(format=vs.RGBS, matrix_in_s=matrix_s).bm3d.RGB2OPP(sample=1)
+
+    def to_fullgray(clip: vs.VideoNode) -> vs.VideoNode:
+        return get_y(clip).resize.Point(format=vs.GRAYS, range_in=Range.LIMITED, range=Range.FULL)
+
+    sigmal = [sigma] * 3 if isinstance(sigma, float) else sigma + [sigma[-1]]*(3-len(sigma))
+    sigmal = [sigmal[0], 0, 0] if isGray else sigmal
+    isGray = True if sigmal[1] == 0 and sigmal[2] == 0 else isGray
+    if len(sigmal) != 3:
+        raise ValueError("bm3d: 'invalid number of sigma parameters supplied'")
+    radiusl = [0, 0] if radius is None else [radius] * 2 if isinstance(radius, int) \
+        else radius + [radius[-1]]*(2-len(radius))
+    if len(radiusl) != 2:
+        raise ValueError("bm3d: 'invalid number or radius parameters supplied'")
+
+    if sigmal[0] == 0 and sigmal[1] == 0 and sigmal[2] == 0:
+        return clip
+
+    pre = pre if pre is None else to_opp(pre) if not isGray else to_fullgray(pre)
+
+    def basic(clip: vs.VideoNode) -> vs.VideoNode:
+        return clip.bm3d.Basic(sigma=sigmal, ref=pre, matrix=100, **basic_args) if radiusl[0] < 1 \
+            else clip.bm3d.VBasic(sigma=sigmal, ref=pre, radius=radiusl[0], matrix=100, **basic_args) \
+            .bm3d.VAggregate(radius=radiusl[0], sample=1)
+
+    clip_in = to_opp(clip) if not isGray else to_fullgray(clip)
+    refv = basic(clip_in) if ref is None else to_opp(ref) if not isGray else to_fullgray(ref)
+
+    def final(clip: vs.VideoNode) -> vs.VideoNode:
+        return clip.bm3d.Final(sigma=sigmal, ref=refv, matrix=100, **final_args) if radiusl[1] < 1 \
+            else clip.bm3d.VFinal(sigma=sigmal, ref=refv, radius=radiusl[1], matrix=100, **final_args) \
+            .bm3d.VAggregate(radius=radiusl[1], sample=1)
+
+    den = iterate(clip_in, final, refine)
+
+    # boil everything back down to whatever input we had
+    den = den.bm3d.OPP2RGB(sample=1).resize.Bicubic(format=clip.format.id, matrix_s=matrix_s) if not isGray \
+        else den.resize.Point(format=clip.format.replace(color_family=vs.GRAY, subsampling_w=0, subsampling_h=0).id,
+                              range_in=Range.FULL, range=Range.LIMITED)
+    # merge source chroma if it exists and we didn't denoise it
+    den = core.std.ShufflePlanes([den, clip], planes=[0, 1, 2], colorfamily=vs.YUV) \
+        if isGray and clip.format.color_family == vs.YUV else den
+    # sub clip luma back in if we only denoised chroma
+    den = den if sigmal[0] != 0 else core.std.ShufflePlanes([clip, den], planes=[0, 1, 2], colorfamily=vs.YUV)
+
+    return den