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
|
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 = {
|
||||||
|
|
Loading…
Reference in New Issue