Skip to content

Instantly share code, notes, and snippets.

@4re
Last active November 3, 2020 17:05
Show Gist options
  • Save 4re/4788aa29cee72ac78361 to your computer and use it in GitHub Desktop.
Save 4re/4788aa29cee72ac78361 to your computer and use it in GitHub Desktop.
from vapoursynth import core, GRAY, YUV
__version__ = "1.0.3"
def knlm_yuv(clip, d=[None, None, None], a=[None, None, None], s=[None, None, None], cmode=None, wmode=None,
h=[None, None, None], rclip=None, device_type=None, info=None, planes=[0]):
"""
Enhanced version of knlmeanscl that supports filtering of all planes on the same call and
pads the image so borders don't be left unfiltered.
Args:
<Default knlmeanscl args>: See knlmeanscl for help.
planes: Default to "0", so only luma is filtered. Set to "[0, 1, 2]" to filter
luma and chroma.
"""
d = populate(d)
a = populate(a)
s = populate(s)
h = populate(h)
if not isinstance(planes, list):
planes = [planes]
packed_y = dict(d=d[0], a=a[0], s=s[0], cmode=cmode, wmode=wmode, h=h[0], rclip=rclip,
device_type=device_type, info=info)
packed_u = dict(d=d[1], a=a[1], s=s[1], cmode=cmode, wmode=wmode, h=h[1], rclip=rclip,
device_type=device_type, info=info)
packed_v = dict(d=d[2], a=a[2], s=s[2], cmode=cmode, wmode=wmode, h=h[2], rclip=rclip,
device_type=device_type, info=info)
pad = max(a)
ret = padding(clip, pad, pad, pad, pad)
y = core.std.ShufflePlanes(ret, planes=[0], colorfamily=GRAY)
u = core.std.ShufflePlanes(ret, planes=[1], colorfamily=GRAY)
v = core.std.ShufflePlanes(ret, planes=[2], colorfamily=GRAY)
if 0 in planes:
y = core.knlm.KNLMeansCL(y, **packed_y)
if 1 in planes:
u = core.knlm.KNLMeansCL(u, **packed_u)
if 2 in planes:
v = core.knlm.KNLMeansCL(v, **packed_v)
ret = core.std.ShufflePlanes([y, u, v], planes=[0, 0, 0], colorfamily=YUV)
ret = core.std.CropRel(ret, pad, pad, pad, pad)
return ret
def populate(var):
"""
Converts any variable in a list of three elements.
The previous value is copied in the subsecuent values.
"""
if not isinstance(var, list):
var = [var] * 3
if len(var) == 1:
var.append(var[0])
if len(var) == 2:
var.append(var[1])
return var
def padding(src, left, top, right, bottom):
"""
Adds specified padding to a video.
"""
w = src.width
h = src.height
last = core.fmtc.resample(src, w=w+left+right, h=h+top+bottom,
sx=-left, sy=-top, sw=w+left+right, sh=h+top+bottom,
kernel="point")
if last.format.bits_per_sample != src.format.bits_per_sample:
last = core.fmtc.bitdepth(last, bits=src.format.bits_per_sample)
return last
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment