This commit is contained in:
comfyanonymous 2023-08-17 03:21:20 -04:00
commit 1498f1a342

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
@ -277,6 +279,35 @@ class FeatherMask:
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,)
NODE_CLASS_MAPPINGS = { NODE_CLASS_MAPPINGS = {
@ -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 = {