add MaskFromRGBCMYBW and MaskSmooth

essentials.py

@@ -655,6 +655,97 @@ class MaskFromSegmentation:
 
         return (mask, )
 
+class MaskFromRGBCMYBW:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "threshold_r": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
+                "threshold_g": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
+                "threshold_b": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
+                "remove_isolated_pixels": ("INT", { "default": 0, "min": 0, "max": 32, "step": 1, }),
+                "fill_holes": ("BOOLEAN", { "default": False }),
+            }
+        }
+
+    RETURN_TYPES = ("MASK","MASK","MASK","MASK","MASK","MASK","MASK","MASK",)
+    RETURN_NAMES = ("red","green","blue","cyan","magenta","yellow","black","white",)
+    FUNCTION = "execute"
+    CATEGORY = "essentials"
+
+    def execute(self, image, threshold_r, threshold_g, threshold_b, remove_isolated_pixels, fill_holes):
+        red = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] < threshold_b)).float()
+        green = ((image[..., 0] < threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] < threshold_b)).float()
+        blue = ((image[..., 0] < threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
+
+        cyan = ((image[..., 0] < threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
+        magenta = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] > 1-threshold_b)).float()
+        yellow = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] < threshold_b)).float()
+
+        black = ((image[..., 0] <= threshold_r) & (image[..., 1] <= threshold_g) & (image[..., 2] <= threshold_b)).float()
+        white = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
+
+        if remove_isolated_pixels > 0 or fill_holes:
+            colors = [red, green, blue, cyan, magenta, yellow, black, white]
+            color_names = ['red', 'green', 'blue', 'cyan', 'magenta', 'yellow', 'black', 'white']
+            processed_colors = {}
+
+            for color_name, color in zip(color_names, colors):
+                color = color.cpu().numpy()
+                masks = []
+
+                for i in range(image.shape[0]):
+                    mask = color[i]
+                    if remove_isolated_pixels > 0:
+                        mask = scipy.ndimage.binary_opening(mask, structure=np.ones((remove_isolated_pixels, remove_isolated_pixels)))
+                    if fill_holes:
+                        mask = scipy.ndimage.binary_fill_holes(mask)
+                    mask = torch.from_numpy(mask)
+                    masks.append(mask)
+
+                processed_colors[color_name] = torch.stack(masks, dim=0).float()
+
+            red = processed_colors['red']
+            green = processed_colors['green']
+            blue = processed_colors['blue']
+            cyan = processed_colors['cyan']
+            magenta = processed_colors['magenta']
+            yellow = processed_colors['yellow']
+            black = processed_colors['black']
+            white = processed_colors['white']
+
+            del colors, processed_colors
+        
+        return (red, green, blue, cyan, magenta, yellow, black, white,)
+
+class MaskSmooth:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "amount": ("INT", { "default": 0, "min": 0, "max": 127, "step": 1, }),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "execute"
+    CATEGORY = "essentials"
+
+    def execute(self, mask, amount):
+        if amount == 0:
+            return (mask,)
+        
+        if amount % 2 == 0:
+            amount += 1
+
+        mask = mask > 0.5
+        mask = T.functional.gaussian_blur(mask.unsqueeze(1), amount).squeeze(1).float()
+
+        return (mask,)
+
+
 class MaskFromBatch:
     @classmethod
     def INPUT_TYPES(s):
@@ -1740,6 +1831,8 @@ NODE_CLASS_MAPPINGS = {
     "MaskFromBatch+": MaskFromBatch,
     "MaskBoundingBox+": MaskBoundingBox,
     "MaskFromSegmentation+": MaskFromSegmentation,
+    "MaskFromRGBCMYBW+": MaskFromRGBCMYBW,
+    "MaskSmooth+": MaskSmooth,
 
     "SimpleMath+": SimpleMath,
     "ConsoleDebug+": ConsoleDebug,
@@ -1791,6 +1884,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "MaskFromBatch+": "🔧 Mask From Batch",
     "MaskBoundingBox+": "🔧 Mask Bounding Box",
     "MaskFromSegmentation+": "🔧 Mask From Segmentation",
+    "MaskFromRGBCMYBW+": "🔧 Mask From RGB/CMY/BW",
+    "MaskSmooth+": "🔧 Mask Smooth",
 
     "SimpleMath+": "🔧 Simple Math",
     "ConsoleDebug+": "🔧 Console Debug",