add a node to allow growing of masks through dilation

This commit is contained in:
Corey 2023-08-16 10:57:14 -04:00
parent 27b87c25a1
commit 18e86a4010
1 changed files with 32 additions and 0 deletions

View File

@ -1,3 +1,5 @@
import numpy as np
from scipy.ndimage import grey_dilation
import torch import torch
from nodes import MAX_RESOLUTION from nodes import MAX_RESOLUTION
@ -276,6 +278,35 @@ class FeatherMask:
output[-y, :] *= feather_rate output[-y, :] *= feather_rate
return (output,) return (output,)
class GrowMask:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mask": ("MASK",),
"expand": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
"tapered_corners": ("BOOLEAN", {"default": True}),
},
}
CATEGORY = "mask"
RETURN_TYPES = ("MASK",)
FUNCTION = "expand_mask"
def expand_mask(self, mask, expand, tapered_corners):
c = 0 if tapered_corners else 1
kernel = np.array([[c, 1, c],
[1, 1, 1],
[c, 1, c]])
output = mask.numpy().copy()
while expand > 0:
output = grey_dilation(output, footprint=kernel)
expand -= 1
output = torch.from_numpy(output)
return (output,)
@ -289,6 +320,7 @@ NODE_CLASS_MAPPINGS = {
"CropMask": CropMask, "CropMask": CropMask,
"MaskComposite": MaskComposite, "MaskComposite": MaskComposite,
"FeatherMask": FeatherMask, "FeatherMask": FeatherMask,
"GrowMask": GrowMask,
} }
NODE_DISPLAY_NAME_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = {