add a node to allow growing of masks through dilation
This commit is contained in:
parent
27b87c25a1
commit
18e86a4010
|
@ -1,3 +1,5 @@
|
|||
import numpy as np
|
||||
from scipy.ndimage import grey_dilation
|
||||
import torch
|
||||
|
||||
from nodes import MAX_RESOLUTION
|
||||
|
@ -277,6 +279,35 @@ class FeatherMask:
|
|||
|
||||
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 = {
|
||||
|
@ -289,6 +320,7 @@ NODE_CLASS_MAPPINGS = {
|
|||
"CropMask": CropMask,
|
||||
"MaskComposite": MaskComposite,
|
||||
"FeatherMask": FeatherMask,
|
||||
"GrowMask": GrowMask,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
|
|
Loading…
Reference in New Issue