ColorToMask

Documentation

• Class name: ColorToMask
• Category: KJNodes/masking
• Output node: False

The ColorToMask node is designed to convert a specified RGB color in an image to a mask. This process involves analyzing the image to identify pixels that match the given RGB values and creating a mask based on these pixels. It supports batch processing, allowing multiple images to be processed simultaneously with control over the batch size.

Input types

Required

• images
  • The images to be processed. This parameter is crucial for defining the input images on which the RGB to mask conversion will be performed.
  • Comfy dtype: IMAGE
  • Python dtype: torch.Tensor
• invert
  • A boolean flag to invert the mask. When set to True, the areas not matching the specified RGB color will be masked instead.
  • Comfy dtype: BOOLEAN
  • Python dtype: bool
• red
  • The red component of the RGB color to be converted to a mask.
  • Comfy dtype: INT
  • Python dtype: int
• green
  • The green component of the RGB color to be converted to a mask.
  • Comfy dtype: INT
  • Python dtype: int
• blue
  • The blue component of the RGB color to be converted to a mask.
  • Comfy dtype: INT
  • Python dtype: int
• threshold
  • The threshold for color matching. Pixels within this threshold of the specified RGB color will be included in the mask.
  • Comfy dtype: INT
  • Python dtype: int
• per_batch
  • Controls the number of images processed in a single batch, allowing for efficient batch processing of images.
  • Comfy dtype: INT
  • Python dtype: int

Output types

• mask
  • Comfy dtype: MASK
  • The output mask generated by identifying pixels that match the specified RGB color within the threshold.
  • Python dtype: torch.Tensor

Usage tips

• Infra type: GPU
• Common nodes: unknown

Source code

class ColorToMask:

    RETURN_TYPES = ("MASK",)
    FUNCTION = "clip"
    CATEGORY = "KJNodes/masking"
    DESCRIPTION = """
Converts chosen RGB value to a mask.  
With batch inputs, the **per_batch**  
controls the number of images processed at once.
"""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                 "images": ("IMAGE",),
                 "invert": ("BOOLEAN", {"default": False}),
                 "red": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}),
                 "green": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}),
                 "blue": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}),
                 "threshold": ("INT", {"default": 10,"min": 0, "max": 255, "step": 1}),
                 "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}),
        },
    } 

    def clip(self, images, red, green, blue, threshold, invert, per_batch):

        color = torch.tensor([red, green, blue], dtype=torch.uint8)  
        black = torch.tensor([0, 0, 0], dtype=torch.uint8)
        white = torch.tensor([255, 255, 255], dtype=torch.uint8)

        if invert:
            black, white = white, black

        steps = images.shape[0]
        pbar = comfy.utils.ProgressBar(steps)
        tensors_out = []

        for start_idx in range(0, images.shape[0], per_batch):

            # Calculate color distances
            color_distances = torch.norm(images[start_idx:start_idx+per_batch] * 255 - color, dim=-1)

            # Create a mask based on the threshold
            mask = color_distances <= threshold

            # Apply the mask to create new images
            mask_out = torch.where(mask.unsqueeze(-1), white, black).float()
            mask_out = mask_out.mean(dim=-1)

            tensors_out.append(mask_out.cpu())
            batch_count = mask_out.shape[0]
            pbar.update(batch_count)

        tensors_out = torch.cat(tensors_out, dim=0)
        tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0)
        return tensors_out,