🔧 Mask From Segmentation

Documentation

Class name: MaskFromSegmentation+
Category: essentials
Output node: False

The node is designed to generate a mask based on image segmentation, facilitating the extraction of specific features or objects from an image by identifying and isolating them through segmentation techniques.

Input types

Required

image
- The input image for segmentation, serving as the basis for mask generation.
- Comfy dtype: IMAGE
- Python dtype: torch.Tensor
segments
- Specifies the segments to be used for mask creation, allowing for targeted extraction within the image.
- Comfy dtype: INT
- Python dtype: List[int]
remove_isolated_pixels
- A boolean flag to remove isolated pixels in the generated mask, enhancing mask quality by eliminating noise.
- Comfy dtype: INT
- Python dtype: bool
remove_small_masks
- A boolean flag to remove small masks from the generated mask, focusing on significant segments by filtering out minor ones.
- Comfy dtype: FLOAT
- Python dtype: bool
fill_holes
- A boolean flag to fill holes in the generated mask, ensuring a more continuous and coherent mask output.
- Comfy dtype: BOOLEAN
- Python dtype: bool

Output types

mask
- Comfy dtype: MASK
- The output mask generated from the specified image segments, isolating desired features or objects.
- Python dtype: torch.Tensor

Usage tips

Infra type: CPU
Common nodes: unknown

Source code

class MaskFromSegmentation:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", ),
                "segments": ("INT", { "default": 6, "min": 1, "max": 16, "step": 1, }),
                "remove_isolated_pixels": ("INT", { "default": 0, "min": 0, "max": 32, "step": 1, }),
                "remove_small_masks": ("FLOAT", { "default": 0.0, "min": 0., "max": 1., "step": 0.01, }),
                "fill_holes": ("BOOLEAN", { "default": False }),
            }
        }

    RETURN_TYPES = ("MASK",)
    FUNCTION = "execute"
    CATEGORY = "essentials"

    def execute(self, image, segments, remove_isolated_pixels, fill_holes, remove_small_masks):
        im = image[0] # we only work on the first image in the batch
        im = Image.fromarray((im * 255).to(torch.uint8).cpu().numpy(), mode="RGB")
        im = im.quantize(palette=im.quantize(colors=segments), dither=Image.Dither.NONE)       
        im = torch.tensor(np.array(im.convert("RGB"))).float() / 255.0

        colors = im.reshape(-1, im.shape[-1])
        colors = torch.unique(colors, dim=0)

        masks = []
        for color in colors:
            mask = (im == color).all(dim=-1).float()
            # remove isolated pixels
            if remove_isolated_pixels > 0:
                mask_np = mask.cpu().numpy()
                mask_np = scipy.ndimage.binary_opening(mask_np, structure=np.ones((remove_isolated_pixels, remove_isolated_pixels)))
                mask = torch.from_numpy(mask_np)

            # fill holes
            if fill_holes:
                mask_np = mask.cpu().numpy()
                mask_np = scipy.ndimage.binary_fill_holes(mask_np)
                mask = torch.from_numpy(mask_np)

            # if the mask is too small, it's probably noise
            if mask.sum() / (mask.shape[0]*mask.shape[1]) > remove_small_masks:
                masks.append(mask)

        if masks == []:
            masks.append(torch.zeros_like(im).squeeze(-1).unsqueeze(0)) # return an empty mask if no masks were found, prevents errors

        mask = torch.stack(masks, dim=0).float()

        return (mask, )

3.7 KiB Raw Permalink Blame History