From 534a9f57e748ea57d8bb49205cc4fa8274f5db20 Mon Sep 17 00:00:00 2001 From: emotion3459 <176516814+emotion3459@users.noreply.github.com> Date: Sun, 15 Dec 2024 12:24:56 -0500 Subject: [PATCH] decouple compensate/degrain tr --- vsdenoise/mvtools/mvtools.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/vsdenoise/mvtools/mvtools.py b/vsdenoise/mvtools/mvtools.py index 1f01af2..be44913 100644 --- a/vsdenoise/mvtools/mvtools.py +++ b/vsdenoise/mvtools/mvtools.py @@ -532,7 +532,8 @@ def compensate( # type: ignore self, func: Union[ Callable[Concatenate[vs.VideoNode, P], vs.VideoNode], Callable[Concatenate[list[vs.VideoNode], P], vs.VideoNode] - ], thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), + ] | None = None, + tr: int | None = None, thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), supers: SuperClips | None = None, *args: P.args, ref: vs.VideoNode | None = None, **kwargs: P.kwargs ) -> vs.VideoNode: @@ -579,8 +580,8 @@ def compensate( # type: ignore @overload def compensate( - self, func: None, - thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), + self, func: None = None, + tr: int | None = None, thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), supers: SuperClips | None = None, ref: vs.VideoNode | None = None ) -> tuple[vs.VideoNode, tuple[int, int]]: """ @@ -626,16 +627,18 @@ def compensate( # type: ignore self, func: Union[ Callable[Concatenate[vs.VideoNode, P], vs.VideoNode], Callable[Concatenate[list[vs.VideoNode], P], vs.VideoNode] - ] | None, thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), + ] | None = None, + tr: int | None = None, thSAD: int = 150, thSCD: int | tuple[int | None, int | None] | None = (None, 51), supers: SuperClips | None = None, *args: P.args, ref: vs.VideoNode | None = None, **kwargs: P.kwargs ) -> vs.VideoNode | tuple[vs.VideoNode, tuple[int, int]]: ref = self.get_ref_clip(ref, self.compensate) + tr = min(tr, self.tr) if tr else self.tr thSCD1, thSCD2 = self.normalize_thscd(thSCD, thSAD, self.compensate) supers = supers or self.get_supers(ref, inplace=True) - vect_b, vect_f = self.get_vectors_bf(self.vectors) + vect_b, vect_f = self.get_vectors_bf(self.vectors, tr=tr) compensate_args = dict( super=supers.render, thsad=thSAD, @@ -646,10 +649,10 @@ def compensate( # type: ignore comp_back, comp_forw = [ [self.mvtools.Compensate(ref, vectors=vect, **compensate_args) for vect in vectors] - for vectors in (reversed(vect_b), vect_f) + for vectors in (vect_b, vect_f) ] - comp_clips = [*comp_forw, ref, *comp_back] + comp_clips = [*comp_back, ref, *comp_forw] n_clips = len(comp_clips) offset = (n_clips - 1) // 2 @@ -732,6 +735,7 @@ def flow( # type: ignore def degrain( self, + tr: int | None = None, thSAD: int | tuple[int | None, int | None] | None = None, limit: int | tuple[int, int] = 255, thSCD: int | tuple[int | None, int | None] | None = (None, 51), @@ -779,13 +783,14 @@ def degrain( """ ref = self.get_ref_clip(ref, self.degrain) + tr = min(tr, self.tr) if tr else self.tr if isinstance(vectors, MVTools): vectors = vectors.vectors elif vectors is None: vectors = self.vectors - vect_b, vect_f = self.get_vectors_bf(vectors, supers=supers, ref=ref) + vect_b, vect_f = self.get_vectors_bf(vectors, supers=supers, ref=ref, tr=tr) supers = supers or self.get_supers(ref, inplace=True) thSAD, thSADC = (thSAD if isinstance(thSAD, tuple) else (thSAD, None)) @@ -810,7 +815,7 @@ def degrain( if self.mvtools is MVToolsPlugin.FLOAT_NEW: output = self.mvtools.Degrain()(ref, supers.render, vectors.vmulti, **degrain_args) else: - output = self.mvtools.Degrain(self.tr)( + output = self.mvtools.Degrain(tr)( ref, supers.render, *chain.from_iterable(zip(vect_b, vect_f)), **degrain_args ) @@ -937,7 +942,7 @@ def get_supers(self, ref: vs.VideoNode, *, inplace: bool = False) -> SuperClips: def get_vectors_bf( self, vectors: MotionVectors, *, supers: SuperClips | None = None, - ref: vs.VideoNode | None = None, inplace: bool = False + ref: vs.VideoNode | None = None, tr: int | None = None, inplace: bool = False ) -> tuple[list[vs.VideoNode], list[vs.VideoNode]]: """ Get the backwards and forward vectors.\n @@ -954,7 +959,8 @@ def get_vectors_bf( if not vectors.has_vectors: vectors = self.analyze(supers=supers, ref=ref, inplace=inplace) - t2 = (self.tr * 2 if self.tr > 1 else self.tr) if self.source_type.is_inter else self.tr + tr = min(tr, self.tr) if tr else self.tr + t2 = (tr * 2 if tr > 1 else tr) if self.source_type.is_inter else tr vectors_backward = list[vs.VideoNode]() vectors_forward = list[vs.VideoNode]()