DeepFaceLive/xlib/avecl/_internal/op/binary_morph.py
2021-10-06 14:50:53 +04:00

46 lines
1.5 KiB
Python

from ..Tensor import Tensor
from .binary_dilate_circle import binary_dilate_circle
from .binary_erode_circle import binary_erode_circle
from .gaussian_blur import gaussian_blur
from .pad import pad
from .cast import cast
def binary_morph(input_t : Tensor, erode_dilate : int, blur : float, fade_to_border : bool = False, dtype=None) -> Tensor:
"""
Apply optional binary erode/dilate and optional blur.
input_t (...,H,W) tensor. Non zero values will be treated as 1.
erode_dilate int >= 0 amount of pixels to dilate
blur float >= 0 amount of pixels to blur
fade_to_border(False) clip the image in order
to fade smoothly to the border with specified blur amount
"""
x = input_t
H,W = input_t.shape[-2:]
x = pad(x, (...,(H,H),(W,W)), mode='constant', constant_value=0)
if erode_dilate > 0:
x = binary_erode_circle(x, radius=1, iterations=max(1,erode_dilate//2))
elif erode_dilate < 0:
x = binary_dilate_circle(x, radius=1, iterations=max(1,-erode_dilate//2) )
if fade_to_border:
h_clip_size = H + blur // 2
w_clip_size = W + blur // 2
x[...,:h_clip_size,:] = 0
x[...,-h_clip_size:,:] = 0
x[...,:,:w_clip_size] = 0
x[...,:,-w_clip_size:] = 0
if blur > 0:
x = gaussian_blur(x, blur * 0.250, dtype=dtype)
else:
x = cast(x, dtype=dtype)
return x[...,H:-H,W:-W]