From 5c0d562b1d3cd81893f8614dc68ade37163fe0cf Mon Sep 17 00:00:00 2001 From: emotion3459 <176516814+emotion3459@users.noreply.github.com> Date: Thu, 14 Nov 2024 18:48:06 -0500 Subject: [PATCH] rewrite vine --- vsdehalo/__init__.py | 2 +- vsdehalo/denoise.py | 69 +++++++++++-- vsdehalo/vine.py | 235 ------------------------------------------- 3 files changed, 64 insertions(+), 242 deletions(-) delete mode 100644 vsdehalo/vine.py diff --git a/vsdehalo/__init__.py b/vsdehalo/__init__.py index 85eca47..dbeaa8b 100644 --- a/vsdehalo/__init__.py +++ b/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 * diff --git a/vsdehalo/denoise.py b/vsdehalo/denoise.py index fe4eccd..a905622 100644 --- a/vsdehalo/denoise.py +++ b/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) \ No newline at end of file diff --git a/vsdehalo/vine.py b/vsdehalo/vine.py deleted file mode 100644 index 71bb061..0000000 --- a/vsdehalo/vine.py +++ /dev/null @@ -1,235 +0,0 @@ -""" -Modern rewrite of IFeelBloated's package with scene detection. -""" -from __future__ import annotations - -from functools import partial -from math import log -from typing import Any - -from vsdenoise import ( - CCDMode, CCDPoints, MotionMode, MVTools, PelType, Prefilter, SearchMode, ccd, frequency_merge, nl_means -) -from vsexprtools import ExprOp, norm_expr_planes -from vskernels import Bicubic, Gaussian -from vsmasktools import Morpho -from vsrgtools import contrasharpening_dehalo, gauss_blur -from vstools import ( - CustomIndexError, CustomRuntimeError, MatrixT, PlanesT, Colorspace, check_ref_clip, check_variable, core, get_y, - join, normalize_planes, split, vs -) - -from .alpha import fine_dehalo - -__all__ = [ - 'super_clip', 'smooth_clip', 'dehalo' -] - - -def super_clip(src: vs.VideoNode, pel: int = 1, planes: PlanesT = 0) -> vs.VideoNode: - assert check_variable(src, super_clip) - - if planes == [1, 2] and (src.format.subsampling_w or src.format.subsampling_h): - src = src.resize.Bicubic( - src.width // (1 << src.format.subsampling_w), - src.height // (1 << src.format.subsampling_h), - format=src.format.replace(subsampling_w=0, subsampling_h=0).id - ) - elif planes == [0]: - src = get_y(src) - - return PelType.NNEDI3(src, pel, nns=4, qual=2, etype=1) - - -def smooth_clip( - src: vs.VideoNode, sr: int = 32, strength: float = 12.5, - sharp: float = 0.80, cutoff: float = 4, - aggressive: bool = True, fast: bool = False, - matrix: MatrixT | None = None, - pel_type: PelType = PelType.BICUBIC, - planes: PlanesT = 0 -) -> vs.VideoNode: - assert check_variable(src, smooth_clip) - - resampler = Colorspace.OPP_BM3D.resampler - - if sr < 1: - raise CustomIndexError('"sr" has to be greater than 0!', smooth_clip) - - if strength <= 0: - raise CustomIndexError('"strength" has to be greater than 0!', smooth_clip) - - if sharp <= 0.0: - raise CustomIndexError('"sharp" has to be greater than 0!', smooth_clip) - - if cutoff < 1 or cutoff > 100: - raise CustomIndexError('"cutoff" must fall in [1, 100]!', smooth_clip) - - csp = src.format.color_family - planes = normalize_planes(src, planes) - - if fast and csp == vs.GRAY: - raise CustomRuntimeError('fast=True is available only for YUV and RGB input!', smooth_clip) - - work_clip = src - - if csp == vs.RGB: - work_clip = resampler.rgb2csp(work_clip, True) # type: ignore - - work_clip, *chroma = split(work_clip) if planes == [0] and not fast else (work_clip, ) # type: ignore - - assert work_clip.format - - # joy - c1 = 0.3926327792690057290863679493724 * sharp - c2 = 18.880334973195822973214959957208 - c3 = 0.5862453661304626725671053478676 - - weight = c1 * log(1.0 + 1.0 / (c1)) - - h_smoothine = c2 * pow(strength / c2, c3) - - upsampled = pel_type(work_clip, 2) - - if fast: - upsampled = ccd( - upsampled, strength * 2, 0, None, - CCDMode.BICUBIC_LUMA, 25 / (2 * sr + 1), matrix, - CCDPoints.MEDIUM, False, planes - ) - upsampled = gauss_blur(upsampled, 1.45454, planes=planes) - else: - upsampled = nl_means(upsampled, strength, 0, sr, 0, planes=planes) - - c = 0 if pel_type == PelType.NNEDI3 else (1.0 - abs(sharp)) / 2 - - resampled = Bicubic(b=-sharp, c=c).scale(upsampled, work_clip.width, work_clip.height) - - if fast: - resampled = contrasharpening_dehalo(resampled, work_clip, 2.5, planes=planes) - - clean = nl_means(work_clip, strength, 0, sr, 0, planes=planes) - - clean = resampled.std.Merge( - clean, [weight if i in planes else 0 for i in range(work_clip.format.num_planes)] - ) - - blur_func = partial(gauss_blur, sigma=Gaussian.sigma.from_fmtc(cutoff), planes=planes) - - if aggressive: - clean = core.std.Expr( - [blur_func(work_clip), clean, blur_func(clean)], - norm_expr_planes(work_clip, 'y z - x +', planes) - ) - else: - clean = frequency_merge( - gauss_blur(work_clip, 0.5), clean, lowpass=blur_func, planes=planes # type: ignore - ) - - diff = work_clip.std.MakeDiff(clean, planes) - - diff = nl_means(diff, h_smoothine, 0, sr, 1, planes=planes, rclip=clean) - - smooth = clean.std.MergeDiff(diff, planes) - - if chroma: - smooth = join(smooth, *chroma) - - if csp == vs.RGB: - return resampler.csp2rgb(smooth, True) - - return smooth - - -def dehalo( - src: vs.VideoNode, smooth: vs.VideoNode | None = None, - tr: int = 0, refine: int = 3, pel: int = 1, thSAD: int = 400, - super_clips: tuple[vs.VideoNode | None, vs.VideoNode | None] = (None, None), - planes: PlanesT = 0, mask: bool | vs.VideoNode = True -) -> vs.VideoNode: - assert check_variable(src, dehalo) - - resampler = Colorspace.OPP_BM3D.resampler - - if smooth: - assert smooth.format - - if src.format.id != smooth.format.id or (src.width, src.height) != (smooth.width, smooth.height): - raise TypeError('vine.dehalo: smooth must have the same format and size as src!') - - if pel not in {1, 2, 4}: - raise CustomIndexError('"pel" has to be 1, 2 o, r 4!') - - csp = src.format.color_family - planes = normalize_planes(src, planes) - - if any(sclip and sclip.format and sclip.format.num_planes == 1 for sclip in super_clips): - planes = [0] - - if smooth is None: - smooth = smooth_clip(src, planes=planes) - - work_clip, smooth_wclip = src, smooth - - if csp == vs.RGB: - work_clip = resampler.rgb2csp(work_clip, True) # type: ignore - smooth_wclip = resampler.rgb2csp(smooth_wclip, True) - - work_clip, *chroma = split(work_clip) if planes == [0] else (work_clip, ) # type: ignore - smooth_wclip = get_y(smooth_wclip) if planes == [0] else smooth_wclip - - constant = 0.0009948813682897925944723492342 - - me_sad = constant * pow(thSAD, 2.0) * log(1.0 + 1.0 / (constant * thSAD)) - - if isinstance(mask, vs.VideoNode): - check_ref_clip(src, (halo_mask := mask)) # type: ignore - elif mask: - halo_mask = fine_dehalo(work_clip, 2.1, ss=1, edgeproc=0.5, show_mask=True, planes=planes) - halo_mask = Morpho.dilation(halo_mask, 1, planes) - else: - halo_mask = None - - smooth_wclip = smooth_wclip.std.MakeDiff(work_clip, planes) - - src_super, smooth_super = super_clips - - if pel == 1: - if not src_super: - src_super = work_clip - if not smooth_super: - smooth_super = smooth_wclip - - if src_super and smooth_super: - smooth_super = smooth_super.std.MakeDiff(src_super) - - class CustomSubPelClipsMVTools(MVTools): - def get_subpel_clips(self, *args: Any) -> tuple[vs.VideoNode | None, vs.VideoNode | None]: - return (smooth_super, src_super) - - mv = CustomSubPelClipsMVTools( - smooth_wclip, tr, refine, prefilter=Prefilter.NONE, - pel=pel, rfilter=4, planes=planes - ) - - mv.analyze_args |= dict[str, Any](trymany=True, badrange=-24, divide=0) - mv.recalculate_args |= dict[str, Any](smooth=1, divide=0, thsad=me_sad) - - mv.analyze(ref=work_clip, search=SearchMode.EXHAUSTIVE, motion=MotionMode.HIGH_SAD) - - averaged_dif = mv.degrain(thSAD=thSAD) - - averaged_dif = ExprOp.MIN(averaged_dif, smooth_wclip, planes=planes) - - clean = work_clip.std.MergeDiff(averaged_dif, planes) - - if halo_mask: - clean = work_clip.std.MaskedMerge(clean, halo_mask, planes) # type: ignore - - if chroma: - clean = join(clean, *chroma) - - if csp == vs.RGB: - return resampler.csp2rgb(clean, True) - - return clean