diff --git a/vsdenoise/funcs.py b/vsdenoise/funcs.py index abdbbfb..d039a03 100644 --- a/vsdenoise/funcs.py +++ b/vsdenoise/funcs.py @@ -23,7 +23,7 @@ from .mvtools import MotionMode, MVTools, MVToolsPreset, MVToolsPresets, SADMode, SearchMode from .mvtools.enums import SearchModeBase from .mvtools.utils import normalize_thscd -from .nlm import WeightMode, WeightModeAndRef, nl_means +from .nlm import NLMWeightMode, NLMWeightModeAndRef, nl_means from .postprocess import PostProcess, PostProcessConfig from .prefilters import PelType, Prefilter @@ -309,7 +309,7 @@ def schizo_denoise( preset: MVToolsPreset = MVToolsPresets.FAST, prefilter: vs.VideoNode | Prefilter = Prefilter.DFTTEST, smooth: bool | tuple[bool | int, bool | int] = True, - wmode: WeightMode | WeightModeAndRef = WeightMode.WELSCH, + wmode: NLMWeightMode | NLMWeightModeAndRef = NLMWeightMode.WELSCH, contra: int | float | bool = False, aggressive: vs.VideoNode | Prefilter | bool = False, matrix: MatrixT | None = None, range_in: ColorRangeT | None = None, nlm_ref: vs.VideoNode | bool | None = None, **kwargs: Any diff --git a/vsdenoise/nlm.py b/vsdenoise/nlm.py index b1f2818..edffa20 100644 --- a/vsdenoise/nlm.py +++ b/vsdenoise/nlm.py @@ -14,7 +14,7 @@ ) __all__ = [ - 'ChannelMode', 'DeviceType', 'WeightMode', + 'ChannelMode', 'DeviceType', 'NLMWeightMode', 'nl_means' ] @@ -180,7 +180,7 @@ def __call__(self, **kwargs: Any) -> DeviceTypeWithInfo: return DeviceTypeWithInfo(str(self), **kwargs) -class WeightMode(CustomIntEnum): +class NLMWeightMode(CustomIntEnum): WELSCH = 0 """ Welsch weighting function has a faster decay, but still assigns positive weights to dissimilar blocks. @@ -203,18 +203,18 @@ class WeightMode(CustomIntEnum): Modified Bisquare weighting function to be even more robust. """ - def __call__(self, weight_ref: float = 1.0) -> WeightModeAndRef: + def __call__(self, weight_ref: float = 1.0) -> NLMWeightModeAndRef: """ :param weight_ref: Amount of original pixel to contribute to the filter output, relative to the weight of the most similar pixel found. :return: Config with weight mode and ref. """ - return WeightModeAndRef(self, weight_ref) + return NLMWeightModeAndRef(self, weight_ref) -class WeightModeAndRef(NamedTuple): - weight_mode: WeightMode +class NLMWeightModeAndRef(NamedTuple): + weight_mode: NLMWeightMode weight_ref: float @@ -222,7 +222,7 @@ def nl_means( clip: vs.VideoNode, strength: float | Sequence[float] = 1.2, tr: int | Sequence[int] = 1, sr: int | Sequence[int] = 2, simr: int | Sequence[int] = 4, device_type: DeviceType = DeviceType.AUTO, ref: vs.VideoNode | None = None, - wmode: WeightMode | WeightModeAndRef = WeightMode.WELSCH, planes: PlanesT = None, + wmode: NLMWeightMode | NLMWeightModeAndRef = NLMWeightMode.WELSCH, planes: PlanesT = None, **kwargs: Any ) -> vs.VideoNode: """ @@ -260,7 +260,7 @@ def nl_means( nstrength, ntr, nsr, nsimr = to_arr(strength), to_arr(tr), to_arr(sr), to_arr(simr) - wmoder, wref = wmode if isinstance(wmode, WeightModeAndRef) else wmode() + wmoder, wref = wmode if isinstance(wmode, NLMWeightModeAndRef) else wmode() kwargs.update(ref=ref, wmode=wmoder.value, wref=wref)