From 87f93fd4d8ca1be5b4a931700ad8aa14af25510e Mon Sep 17 00:00:00 2001 From: Setsugennoao <41454651+Setsugennoao@users.noreply.github.com> Date: Thu, 14 Dec 2023 16:06:53 +0100 Subject: [PATCH] Add border_handling to ComplexKernel.scale --- vskernels/kernels/complex.py | 42 +++++++++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/vskernels/kernels/complex.py b/vskernels/kernels/complex.py index fb8880d..d020883 100644 --- a/vskernels/kernels/complex.py +++ b/vskernels/kernels/complex.py @@ -1,10 +1,11 @@ from __future__ import annotations +from functools import lru_cache from typing import TYPE_CHECKING, Any, SupportsFloat, TypeVar, Union, cast from stgpytools import inject_kwargs_params from vstools import ( - CustomIntEnum, Dar, KwargsT, Resolution, Sar, VSFunctionAllArgs, check_correct_subsampling, inject_self, vs + CustomIntEnum, Dar, KwargsT, Resolution, Sar, VSFunctionAllArgs, check_correct_subsampling, inject_self, padder, vs ) from ..types import Center, LeftShift, Slope, TopShift @@ -29,6 +30,31 @@ class BorderHandling(CustomIntEnum): ZERO = 1 REPEAT = 2 + def prepare_clip(self, clip: vs.VideoNode, min_pad: int = 2) -> vs.VideoNode: + pad_w, pad_h = ( + self.pad_amount(size, min_pad) for size in (clip.width, clip.height) + ) + + if pad_w == pad_h == 0: + return clip + + args = (clip, pad_w, pad_w, pad_h, pad_h) + + match self: + case BorderHandling.MIRROR: + return padder.MIRROR(*args) + case BorderHandling.ZERO: + return padder.COLOR(*args) + case BorderHandling.REPEAT: + return padder.REPEAT(*args) + + @lru_cache + def pad_amount(self, size: int, min_amount: int = 2) -> int: + if self is BorderHandling.MIRROR: + return 0 + + return (((size + min_amount) + 7) & -8) - size + def _from_param(cls: type[XarT], value: XarT | bool | float | None, fallback: XarT) -> XarT | None: if value is False: @@ -72,7 +98,7 @@ def func( resampler: Resampler | None = self if isinstance(self, Resampler) else None with LinearLight(clip, linear, sigmoid, resampler, kwargs.pop('format', None)) as ll: - ll.linear = operation(ll.linear, width, height, shift, **kwargs) + ll.linear = operation(ll.linear, width, height, shift, **kwargs) # type: ignore return ll.out @@ -174,6 +200,7 @@ def _handle_crop_resize_kwargs( # type: ignore[override] @inject_kwargs_params def scale( # type: ignore[override] self, clip: vs.VideoNode, width: int, height: int, shift: tuple[TopShift, LeftShift] = (0, 0), *, + border_handling: BorderHandling = BorderHandling.MIRROR, sar: Sar | float | bool | None = None, dar: Dar | float | bool | None = None, keep_ar: bool = False, **kwargs: Any ) -> vs.VideoNode: @@ -186,6 +213,12 @@ def scale( # type: ignore[override] kwargs, shift, out_sar = self._handle_crop_resize_kwargs(clip, width, height, shift, **kwargs) + padded = border_handling.prepare_clip(clip, self.kernel_radius) + + shift = tuple( # type: ignore + s + p - c for s, p, c in zip(shift, *((x.width, x.height) for x in (clip, padded))) + ) + kwargs = self.get_scale_args(clip, shift, width, height, **kwargs) clip = self.scale_function(clip, **kwargs) @@ -202,12 +235,15 @@ class ComplexScaler(LinearScaler, KeepArScaler): def scale( # type: ignore[override] self, clip: vs.VideoNode, width: int, height: int, shift: tuple[TopShift, LeftShift] = (0, 0), *, + border_handling: BorderHandling = BorderHandling.MIRROR, sar: Sar | bool | float | None = None, dar: Dar | bool | float | None = None, keep_ar: bool = False, linear: bool = False, sigmoid: bool | tuple[Slope, Center] = False, **kwargs: Any ) -> vs.VideoNode: return super().scale( - clip, width, height, shift, sar=sar, dar=dar, keep_ar=keep_ar, linear=linear, sigmoid=sigmoid, **kwargs + clip, width, height, shift, sar=sar, dar=dar, keep_ar=keep_ar, + linear=linear, sigmoid=sigmoid, border_handling=border_handling, + **kwargs )