Skip to content

Commit

Permalink
Merge masked_clamp_aa into clamp_aa (#26)
Browse files Browse the repository at this point in the history 
* Merge masked_clamp_aa into clamp_aa

* Fix leftover issues

* flake8 fix
  • Loading branch information
@emotion3459
emotion3459 authored Sep 10, 2024
1 parent 6bcd416 commit a32cb05
Showing 1 changed file with 44 additions and 68 deletions.
112 changes: 44 additions & 68 deletions vsaa/funcs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
'pre_aa',
'upscaled_sraa',
'transpose_aa',
'clamp_aa', 'masked_clamp_aa',
'clamp_aa',
'fine_aa',
'based_aa'
]
Expand Down Expand Up @@ -166,54 +166,10 @@ def transpose_aa(clip: vs.VideoNode, aafunc: SingleRater, planes: PlanesT = 0) -


def clamp_aa(
src: vs.VideoNode, weak: vs.VideoNode, strong: vs.VideoNode,
strength: float = 1, ref: vs.VideoNode | None = None, planes: PlanesT = 0
) -> vs.VideoNode:
"""
Clamp stronger AAs to weaker AAs.
Useful for clamping :py:func:`upscaled_sraa` or :py:func:`Eedi3`
to :py:func:`Nnedi3` for a strong but more precise AA.
:param src: Non-AA'd source clip.
:param weak: Weakly-AA'd clip.
:param strong: Strongly-AA'd clip.
:param strength: Clamping strength.
:param ref: Reference clip for clamping.
:param planes: Planes to process.
:return: Clip with clamped anti-aliasing.
"""
assert src.format

planes = normalize_planes(src, planes)

if src.format.sample_type == vs.INTEGER:
thr = strength * get_peak_value(src)
else:
thr = strength / 219

if thr == 0:
return median_clips(src, weak, strong, planes=planes)

if ref:
if complexpr_available:
expr = f'y z - ZD! y a - AD! ZD@ AD@ xor x ZD@ abs AD@ abs < a z {thr} + min z {thr} - max a ? ?'
else:
expr = f'y z - y a - xor x y z - abs y a - abs < a z {thr} + min z {thr} - max a ? ?'
else:
if complexpr_available:
expr = f'x y - XYD! x z - XZD! XYD@ XZD@ xor x XYD@ abs XZD@ abs < z y {thr} + min y {thr} - max z ? ?'
else:
expr = f'x y - x z - xor x x y - abs x z - abs < z y {thr} + min y {thr} - max z ? ?'

return norm_expr([src, ref, weak, strong] if ref else [src, weak, strong], expr, planes)


def masked_clamp_aa(
clip: vs.VideoNode, strength: float = 1.0,
mthr: float = 0.25, mask: vs.VideoNode | EdgeDetectT | None = None,
weak_aa: vs.VideoNode | SingleRater = Nnedi3(),
strong_aa: vs.VideoNode | SingleRater = Eedi3(),
weak_aa: vs.VideoNode | Antialiaser = Nnedi3(),
strong_aa: vs.VideoNode | Antialiaser = Eedi3(),
opencl: bool | None = False, ref: vs.VideoNode | None = None,
planes: PlanesT = 0
) -> vs.VideoNode:
Expand All @@ -224,50 +180,70 @@ def masked_clamp_aa(
:param strength: Set threshold strength for over/underflow value for clamping.
:param mthr: Binarize threshold for the mask, float.
:param mask: Clip to use for custom mask or an EdgeDetect to use custom masker.
:param weak_aa: SingleRater for the weaker aa.
:param strong_aa: SingleRater for the stronger aa.
:param weak_aa: Antialiaser for the weaker aa.
:param strong_aa: Antialiaser for the stronger aa.
:param opencl: Whether to force OpenCL acceleration, None to leave as is.
:param ref: Reference clip for clamping.
:return: Antialiased clip.
"""
assert clip.format

func = FunctionUtil(clip, masked_clamp_aa, planes)

if not isinstance(mask, vs.VideoNode):
bin_thr = scale_value(mthr, 32, clip)
func = FunctionUtil(clip, clamp_aa, planes, vs.YUV)

mask = ScharrTCanny.ensure_obj(mask).edgemask(func.work_clip) # type: ignore
mask = mask.std.Binarize(bin_thr)
mask = mask.std.Maximum()
mask = box_blur(mask)
mask = mask.std.Minimum().std.Deflate()

if isinstance(weak_aa, SingleRater):
if not isinstance(weak_aa, vs.VideoNode):
if opencl is not None and hasattr(weak_aa, 'opencl'):
weak_aa.opencl = opencl

weak_aa = transpose_aa(func.work_clip, weak_aa, func.norm_planes)
weak_aa = weak_aa.aa(func.work_clip)
elif func.luma_only:
weak_aa = get_y(weak_aa)

if isinstance(strong_aa, SingleRater):
if not isinstance(strong_aa, vs.VideoNode):
if opencl is not None and hasattr(strong_aa, 'opencl'):
strong_aa.opencl = opencl

strong_aa = transpose_aa(func.work_clip, strong_aa, func.norm_planes)
strong_aa = strong_aa.aa(func.work_clip)
elif func.luma_only:
strong_aa = get_y(strong_aa)

if ref and func.luma_only:
ref = get_y(ref)

clamped = clamp_aa(func.work_clip, weak_aa, strong_aa, strength, ref, func.norm_planes)
if func.work_clip.format.sample_type == vs.INTEGER:
thr = strength * get_peak_value(func.work_clip)
else:
thr = strength / 219

if thr == 0:
clamped = median_clips([func.work_clip, weak_aa, strong_aa], planes=planes)
else:
if ref:
if complexpr_available:
expr = f'y z - ZD! y a - AD! ZD@ AD@ xor x ZD@ abs AD@ abs < a z {thr} + min z {thr} - max a ? ?'
else:
expr = f'y z - y a - xor x y z - abs y a - abs < a z {thr} + min z {thr} - max a ? ?'
else:
if complexpr_available:
expr = f'x y - XYD! x z - XZD! XYD@ XZD@ xor x XYD@ abs XZD@ abs < z y {thr} + min y {thr} - max z ? ?'
else:
expr = f'x y - x z - xor x x y - abs x z - abs < z y {thr} + min y {thr} - max z ? ?'

clamped = norm_expr(
[func.work_clip, ref, weak_aa, strong_aa] if ref else [func.work_clip, strong_aa, strong_aa],
expr, func.norm_planes
)

if mask:
if not isinstance(mask, vs.VideoNode):
bin_thr = scale_value(mthr, 32, clip)

mask = ScharrTCanny.ensure_obj(mask).edgemask(func.work_clip) # type: ignore
mask = box_blur(mask.std.Binarize(bin_thr).std.Maximum())
mask = mask.std.Minimum().std.Deflate()

merged = func.work_clip.std.MaskedMerge(clamped, mask, func.norm_planes) # type: ignore
clamped = func.work_clip.std.MaskedMerge(clamped, mask, func.norm_planes) # type: ignore

return func.return_clip(merged)
return func.return_clip(clamped)


def fine_aa(
Expand Down Expand Up @@ -343,7 +319,7 @@ def based_aa(
'You\'re missing the "vsdenoise" package! Install it with "pip install vsdenoise".', based_aa
)

func = FunctionUtil(clip, based_aa, planes)
func = FunctionUtil(clip, based_aa, planes, vs.YUV)

if supersampler is False:
supersampler = downscaler = NoScale
Expand Down

0 comments on commit a32cb05

Please sign in to comment.