From 18e86a4010099a39eee8c901e22696575b330f86 Mon Sep 17 00:00:00 2001
From: Corey <17171838+coreyryanhanson@users.noreply.github.com>
Date: Wed, 16 Aug 2023 10:57:14 -0400
Subject: [PATCH] add a node to allow growing of masks through dilation

---
 comfy_extras/nodes_mask.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/comfy_extras/nodes_mask.py b/comfy_extras/nodes_mask.py
index 87ba36fc..5adb468a 100644
--- a/comfy_extras/nodes_mask.py
+++ b/comfy_extras/nodes_mask.py
@@ -1,3 +1,5 @@
+import numpy as np
+from scipy.ndimage import grey_dilation
 import torch
 
 from nodes import MAX_RESOLUTION
@@ -276,6 +278,35 @@ class FeatherMask:
             output[-y, :] *= feather_rate
 
         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,
     "MaskComposite": MaskComposite,
     "FeatherMask": FeatherMask,
+    "GrowMask": GrowMask,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {