From 01945ade869671e725a984f2da615e0be376da66 Mon Sep 17 00:00:00 2001 From: "Salvador E. Tropea" Date: Tue, 16 Sep 2025 13:50:05 -0300 Subject: [PATCH] [Added] Node to apply a mask to an image Using Approximate Fast Foreground Colour Estimation --- README.md | 20 ++++++++++++++++++++ src/nodes/nodes_img.py | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/README.md b/README.md index b5caacc..6489d36 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ Currently we just have a few nodes used by other nodes I maintain. - [Normalize Image to ImageNet](#4-normalize-image-to-imagenet) - [Normalize Image to [-0.5, 0.5]](#5-normalize-image-to-05-05) - [Normalize Image to [-1, 1]](#6-normalize-image-to-1-1) + - [Apply Mask using AFFCE](#7-apply-mask-using-affce) - 📝 [Usage Notes](#-usage-notes) - 📜 [Project History](#-project-history) - ⚖️ [License](#️-license) @@ -128,6 +129,25 @@ Currently we just have a few nodes used by other nodes I maintain. - **Mean:** `[0.5, 0.5, 0.5]` - **Std Dev:** `[0.5, 0.5, 0.5]` +### 7. Apply Mask using AFFCE + +- **Display Name:** `Apply Mask using AFFCE` +- **Internal Name:** `SET_ApplyMaskAFFCE` +- **Category:** `image/manipulation +- **Description:** Applies a mask to an image using [Approximate Fast Foreground Colour Estimation](https://github.com/Photoroom/fast-foreground-estimation). This blends the image contour in a better way. +- **Purpose:** Used to apply the mask of a background removal model. +- **Inputs:** + - `images` (`IMAGE`): One ore more ComfyUI images + - `masks` (`MASK`): Masks to apply + - `blur_size` (`INT`): Diameter for the coarse gaussian blur + - `blur_size_two` (`INT`): Diameter for the fine gaussian blur + - `fill_color` (`BOOLEAN`): When enabled the removed image is replaced by a color, the output is an RGB image. Otherwise the removed part becomes transparent and the output is an RGBA image. + - `color` (`STRING`): A string representing a color to be used when `fill_color` is enabled. Can be an hexadecimal RGB (i.e. `#AABBCC`) or comma separated RGB components. The components can be in the [0-255] or [0-1.0] range. +- **Output:** + - `image` (`IMAGE`): The image after applying the mask. + - `mask` (`MASK`): The input mask + + ## 🚀 Installation You can install the nodes from the ComfyUI nodes manager, the name is *Image Misc*, or just do it manually: diff --git a/src/nodes/nodes_img.py b/src/nodes/nodes_img.py index ae3e861..d5df1ae 100644 --- a/src/nodes/nodes_img.py +++ b/src/nodes/nodes_img.py @@ -7,6 +7,7 @@ import numpy as np import os from PIL import Image # Import the Python Imaging Library +from seconohe.apply_mask import apply_mask from seconohe.downloader import download_file # We are the main source, so we use the main_logger from . import main_logger @@ -15,6 +16,7 @@ import torchvision.transforms.functional as TF from typing import Optional try: from folder_paths import get_input_directory # To get the ComfyUI input directory + from comfy import model_management except ModuleNotFoundError: # No ComfyUI, this is a test environment def get_input_directory(): @@ -31,6 +33,13 @@ BASE_CATEGORY = "image" IO_CATEGORY = "io" MANIPULATION_CATEGORY = "manipulation" NORMALIZATION = "normalization" +BLUR_SIZE_OPT = ("INT", {"default": 90, "min": 1, "max": 255, "step": 1, }) +BLUR_SIZE_TWO_OPT = ("INT", {"default": 6, "min": 1, "max": 255, "step": 1, }) +COLOR_OPT = ("STRING", { + "default": "#000000", + "tooltip": "Color for fill.\n" + "Can be an hexadecimal (#RRGGBB).\n" + "Can comma separated RGB values in [0-255] or [0-1.0] range."}) def tensor_to_pil(tensor: torch.Tensor) -> Image.Image: @@ -387,3 +396,33 @@ class NormalizeToRangeMinus1to1(): return (TF.normalize(image.permute(0, 3, 1, 2), # BHWC -> BCHW mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]).permute(0, 2, 3, 1),) # BCHW -> BHWC + + +class ApplyMaskAFFCE: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "masks": ("MASK",), + "blur_size": BLUR_SIZE_OPT, + "blur_size_two": BLUR_SIZE_TWO_OPT, + "fill_color": ("BOOLEAN", {"default": False}), + "color": COLOR_OPT, + } + } + + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image", "mask",) + FUNCTION = "get_foreground" + CATEGORY = BASE_CATEGORY + "/" + MANIPULATION_CATEGORY + DESCRIPTION = ("Apply a mask to an image using\n" + "Approximate Fast Foreground Colour Estimation.\n" + "https://github.com/Photoroom/fast-foreground-estimation") + UNIQUE_NAME = "SET_ApplyMaskAFFCE" + DISPLAY_NAME = "Apply Mask using AFFCE" + + def get_foreground(self, images, masks, blur_size=91, blur_size_two=7, fill_color=False, color=None): + out_images = apply_mask(logger, images, masks, model_management.get_torch_device(), blur_size, blur_size_two, + fill_color, color) + return out_images.cpu(), masks.cpu()