Add FloodFill node

__init__.py

@@ -1,13 +1,14 @@
 from .py.node.calc_max_font_size import CalcMaxFontSize
 from .py.node.segs_to_region import SegsToRegion
 from .py.node.extract_dominant_color import ExtractDominantColor
-from .py.node.colors import GetComplementaryColor
+from .py.node.colors import GetComplementaryColor, FloodFill
 
 NODE_CLASS_MAPPINGS = {
     "SegsToRegion": SegsToRegion,
     "CalcMaxFontSize": CalcMaxFontSize,
     "ExtractDominantColor": ExtractDominantColor,
     "GetComplementaryColor": GetComplementaryColor,
+    "TextOnSegsFloodFill": FloodFill,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -15,6 +16,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CalcMaxFontSize": "CalcMaxFontSize",
     "ExtractDominantColor": "ExtractDominantColor",
     "GetComplementaryColor": "GetComplementaryColor",
+    "TextOnSegsFloodFill": "FloodFill",
 }
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

py/node/colors.py

@@ -1,4 +1,7 @@
-from PIL import ImageColor
+from PIL import Image, ImageColor, ImageDraw
+
+import numpy as np
+import torch
 
 from .base import BaseNode
 
@@ -37,3 +40,89 @@ class GetComplementaryColor(BaseNode):
             r2, g2, b2 = ImageColor.getcolor(alt_color_code, "RGB")
 
         return (r2, g2, b2, color_code)
+
+
+class FloodFill(BaseNode):
+    RETURN_TYPES = ("IMAGE", "INT", "INT", "INT")
+    RETURN_NAMES = (
+        "image",
+        "filled_color_r",
+        "filled_color_g",
+        "filled_color_b",
+    )
+    FUNCTION = "flood_fill"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
+                "y": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
+                "pick_color": ("BOOLEAN", {"default": True}),
+                "r": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}),
+                "g": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}),
+                "b": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}),
+                "threshold": ("INT", {"default": 50, "min": 0, "max": 1000, "step": 1}),
+            }
+        }
+
+    def flood_fill(self, image, x, y, pick_color, r, g, b, threshold):
+        p = (x, y)
+
+        i = 255.0 * image[0].cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        img = img.convert("RGB")
+        if pick_color:
+            r, g, b = img.getpixel(p)
+        self._flood_fill(img, p, (r, g, b), thresh=threshold)
+
+        filled_image = np.array(img).astype(np.float32) / 255.0
+        filled_image = torch.from_numpy(filled_image)[None,]
+        return filled_image, r, g, b
+
+    # refer. https://github.com/python-pillow/Pillow/blob/5d6f22da122c6620331fd11bc8fd8a88cbc11132/src/PIL/ImageDraw.py#L897
+    def _flood_fill(self, image, xy, value, thresh):
+        pixel = image.load()
+
+        x, y = xy
+        try:
+            background = pixel[x, y]
+            # if _color_diff(value, background) <= thresh:
+            #     return  # seed point already has fill color
+            pixel[x, y] = value
+        except (ValueError, IndexError):
+            return  # seed point outside image
+        edge = {(x, y)}
+        # use a set to keep record of current and previous edge pixels
+        # to reduce memory consumption
+        full_edge = set()
+        while edge:
+            new_edge = set()
+            for x, y in edge:  # 4 adjacent method
+                for s, t in ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)):
+                    # If already processed, or if a coordinate is negative, skip
+                    if (s, t) in full_edge or s < 0 or t < 0:
+                        continue
+                    try:
+                        p = pixel[s, t]
+                    except (ValueError, IndexError):
+                        pass
+                    else:
+                        full_edge.add((s, t))
+                        fill = self._color_diff(p, background) <= thresh
+                        if fill:
+                            pixel[s, t] = value
+                            new_edge.add((s, t))
+            full_edge = edge  # discard pixels processed
+            edge = new_edge
+
+    # refer. https://github.com/python-pillow/Pillow/blob/5d6f22da122c6620331fd11bc8fd8a88cbc11132/src/PIL/ImageDraw.py#L1076
+    def _color_diff(self, color1, color2) -> float:
+        """
+        Uses 1-norm distance to calculate difference between two values.
+        """
+        if isinstance(color2, tuple):
+            return sum(abs(color1[i] - color2[i]) for i in range(0, len(color2)))
+        else:
+            return abs(color1 - color2)