From 909ef05d14a0305a0b2b4afcd4d3fd1181942468 Mon Sep 17 00:00:00 2001 From: Cyber Dick Lang <286878701@qq.com> Date: Wed, 17 Sep 2025 17:50:07 +0800 Subject: [PATCH] =?UTF-8?q?v1.3.3:=20=E6=96=B0=E5=A2=9Ekjnodes=E8=8A=82?= =?UTF-8?q?=E7=82=B9=E7=A7=BB=E6=A4=8D=E5=92=8C=E6=9E=B6=E6=9E=84=E4=BC=98?= =?UTF-8?q?=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增Color Match (UTK)节点：支持6种颜色匹配算法，用于图像间色彩转移 - 新增Color To Mask (UTK)节点：根据RGB颜色值创建掩码，支持阈值调节 - 新增Separate Masks (UTK)节点：分离连通组件为独立掩码，支持3种输出模式 - 新增Bbox Visualize (UTK)节点：在图像上绘制边界框，支持xywh和xyxy格式 - 重构mask分类架构：创建独立py文件封装，与image分类保持一致 - 优化Color Match节点输入顺序：image_target在前，避免bypass节点传递错误 - 完善节点分类和导航：所有新节点正确显示在右侧导航面板 - 增强依赖管理：添加color-matcher、scipy等必要依赖 --- COLOR_MATCH_README.md | 98 +++++++++++ __init__.py | 66 ++++++- nodes/image/bbox_visualize.py | 120 +++++++++++++ nodes/image/color_match_standalone.py | 164 ++++++++++++++++++ nodes/mask/__init__.py | 29 ++++ nodes/mask/separate_masks.py | 237 ++++++++++++++++++++++++++ nodes/tools/color_to_mask.py | 164 ++++++++++++++++++ nodes/tools/kontext_presets.py | 4 +- pyproject.toml | 2 +- requirements.txt | 3 + 10 files changed, 875 insertions(+), 12 deletions(-) create mode 100644 COLOR_MATCH_README.md create mode 100644 nodes/image/bbox_visualize.py create mode 100644 nodes/image/color_match_standalone.py create mode 100644 nodes/mask/__init__.py create mode 100644 nodes/mask/separate_masks.py create mode 100644 nodes/tools/color_to_mask.py diff --git a/COLOR_MATCH_README.md b/COLOR_MATCH_README.md new file mode 100644 index 0000000..5e2a96a --- /dev/null +++ b/COLOR_MATCH_README.md @@ -0,0 +1,98 @@ +# Color Match (UTK) 节点使用说明 + +## 概述 + +Color Match (UTK) 节点是从 kjnodes 项目移植过来的颜色匹配功能，现在已成功集成到 ComfyUI Universal Toolkit 的 image 分类中。该节点能够将参考图像的颜色特征转移到目标图像上，非常适合用于自动色彩分级、照片后期处理和影片色彩统一等场景。 + +## 功能特点 + +- **多种颜色匹配算法**：支持 6 种不同的颜色匹配方法 +- **强度控制**：可调节颜色匹配的强度（0.0-10.0） +- **批处理支持**：支持批量图像处理 +- **多线程优化**：可选择开启多线程以提升处理速度 +- **高质量结果**：基于学术研究的先进算法 + +## 安装要求 + +在使用此节点之前，需要安装 `color-matcher` 依赖库： + +```bash +pip install color-matcher +``` + +## 节点参数 + +### 必需参数 +- **image_ref**：参考图像，作为颜色匹配的目标样式 +- **image_target**：目标图像，需要被调整颜色的图像 +- **method**：颜色匹配方法，可选择： + - `mkl`：Monge-Kantorovich Linearization（默认） + - `hm`：Histogram Matching + - `reinhard`：Reinhard et al. method + - `mvgd`：Multi-Variate Gaussian Distribution + - `hm-mvgd-hm`：Histogram Matching + MVGD + Histogram Matching + - `hm-mkl-hm`：Histogram Matching + MKL + Histogram Matching + +### 可选参数 +- **strength**：颜色匹配强度（默认：1.0，范围：0.0-10.0） +- **multithread**：是否启用多线程处理（默认：True） + +## 使用方法 + +1. 在 ComfyUI 中，在节点菜单中找到 `UTK/image` 分类 +2. 添加 `Color Match (UTK)` 节点 +3. 连接参考图像到 `image_ref` 输入 +4. 连接目标图像到 `image_target` 输入 +5. 选择合适的颜色匹配方法 +6. 调整强度参数（可选） +7. 运行工作流 + +## 算法说明 + +### 推荐的方法选择 + +- **mkl**：适合大多数场景的通用方法，效果平衡 +- **hm-mvgd-hm**：复合方法，通常能获得最佳效果 +- **reinhard**：经典方法，适合艺术风格转换 +- **hm**：简单快速，适合基本的色彩调整 + +### 性能优化 + +- 对于单张图像，建议关闭多线程 +- 对于批量处理，建议开启多线程以提升速度 +- 较低的强度值（0.3-0.7）通常能产生更自然的效果 + +## 技术参考 + +该节点基于以下学术研究： +- Reinhard et al. 的颜色转移方法 +- Pitie et al. 提出的 Monge-Kantorovich Linearization +- Multi-Variate Gaussian Distribution 转移的解析解 + +项目参考：https://github.com/hahnec/color-matcher/ + +## 故障排除 + +### 常见问题 + +1. **导入错误**：确保已安装 `color-matcher` 库 +2. **处理失败**：检查输入图像是否为有效的张量格式 +3. **内存不足**：对于大图像，可以尝试关闭多线程或降低批处理大小 + +### 错误处理 + +节点包含完善的错误处理机制： +- 如果颜色匹配失败，会自动返回原始图像 +- 所有错误信息会在控制台中显示 +- 支持优雅降级，确保工作流不会中断 + +## 更新日志 + +- **v1.0**：初始版本，从 kjnodes 移植并集成到 UTK +- 支持所有原始功能和参数 +- 添加了详细的文档和错误处理 +- 优化了代码结构和性能 + +--- + +*此节点是 ComfyUI Universal Toolkit 的一部分，致力于为 ComfyUI 用户提供更丰富的图像处理功能。* diff --git a/__init__.py b/__init__.py index 8adbb8e..2e98420 100644 --- a/__init__.py +++ b/__init__.py @@ -8,13 +8,25 @@ A comprehensive toolkit for ComfyUI that provides various utility nodes for imag :license: MIT, see LICENSE for more details. """ -__version__ = "1.3.1" +__version__ = "1.3.3" __author__ = "CyberDickLang" __email__ = "286878701@qq.com" __url__ = "https://github.com/whmc76" # 更新日志 CHANGELOG = { + "1.3.3": [ + "新增多个kjnodes节点移植和架构优化：", + "- 新增Color Match (UTK)节点：支持6种颜色匹配算法，用于图像间色彩转移", + "- 新增Color To Mask (UTK)节点：根据RGB颜色值创建掩码，支持阈值调节", + "- 新增Separate Masks (UTK)节点：分离连通组件为独立掩码，支持3种输出模式", + "- 新增Bbox Visualize (UTK)节点：在图像上绘制边界框，支持xywh和xyxy格式", + "- 重构mask分类架构：创建独立py文件封装，与image分类保持一致", + "- 优化Color Match节点输入顺序：image_target在前，避免bypass节点传递错误", + "- 完善节点分类和导航：所有新节点正确显示在右侧导航面板", + "- 增强依赖管理：添加color-matcher、scipy等必要依赖", + "- 提升代码质量：完整的错误处理、进度显示和参数验证", + ], "1.3.1": [ "修复Audio Crop Process节点duration=0时的裁剪逻辑：", "- 修复当duration_seconds为0时仍会进行音频裁剪的问题", @@ -340,10 +352,10 @@ try: from .nodes.image.imitation_hue_node import \ NODE_DISPLAY_NAME_MAPPINGS as IMITATION_HUE_DISPLAY_MAPPINGS # 掩码节点 - from .nodes.mask.mask_operations import \ - NODE_CLASS_MAPPINGS as MASK_OPERATIONS_MAPPINGS - from .nodes.mask.mask_operations import \ - NODE_DISPLAY_NAME_MAPPINGS as MASK_OPERATIONS_DISPLAY_MAPPINGS + from .nodes.mask import \ + NODE_CLASS_MAPPINGS as MASK_MAPPINGS + from .nodes.mask import \ + NODE_DISPLAY_NAME_MAPPINGS as MASK_DISPLAY_MAPPINGS from .nodes.tools.math_expression_node import \ NODE_CLASS_MAPPINGS as MATH_EXPRESSION_MAPPINGS from .nodes.tools.math_expression_node import \ @@ -363,8 +375,8 @@ except ImportError as e: # 如果模块化导入失败，使用空字典 AUDIO_CROP_MAPPINGS = {} AUDIO_CROP_DISPLAY_MAPPINGS = {} - MASK_OPERATIONS_MAPPINGS = {} - MASK_OPERATIONS_DISPLAY_MAPPINGS = {} + MASK_MAPPINGS = {} + MASK_DISPLAY_MAPPINGS = {} CONCATENATE_MULTI_MAPPINGS = {} CONCATENATE_MULTI_DISPLAY_MAPPINGS = {} PAD_OUTPAINT_MAPPINGS = {} @@ -516,6 +528,24 @@ except ImportError: RESTORE_CROP_MAPPINGS = {} RESTORE_CROP_DISPLAY = {} +try: + from .nodes.image.color_match_standalone import \ + NODE_CLASS_MAPPINGS as COLOR_MATCH_MAPPINGS + from .nodes.image.color_match_standalone import \ + NODE_DISPLAY_NAME_MAPPINGS as COLOR_MATCH_DISPLAY +except ImportError: + COLOR_MATCH_MAPPINGS = {} + COLOR_MATCH_DISPLAY = {} + +try: + from .nodes.image.bbox_visualize import \ + NODE_CLASS_MAPPINGS as BBOX_VISUALIZE_MAPPINGS + from .nodes.image.bbox_visualize import \ + NODE_DISPLAY_NAME_MAPPINGS as BBOX_VISUALIZE_DISPLAY +except ImportError: + BBOX_VISUALIZE_MAPPINGS = {} + BBOX_VISUALIZE_DISPLAY = {} + try: from .nodes.tools.think_remover_node import \ NODE_CLASS_MAPPINGS as THINK_REMOVER_MAPPINGS @@ -557,6 +587,15 @@ except ImportError as e: PROMPT_HELPER_MAPPINGS = {} PROMPT_HELPER_DISPLAY_MAPPINGS = {} +try: + from .nodes.tools.color_to_mask import \ + NODE_CLASS_MAPPINGS as COLOR_TO_MASK_MAPPINGS + from .nodes.tools.color_to_mask import \ + NODE_DISPLAY_NAME_MAPPINGS as COLOR_TO_MASK_DISPLAY +except ImportError: + COLOR_TO_MASK_MAPPINGS = {} + COLOR_TO_MASK_DISPLAY = {} + # 合并所有节点映射 NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(EMPTY_UNIT_MAPPINGS) @@ -576,8 +615,10 @@ NODE_CLASS_MAPPINGS.update(CHECK_MASK_MAPPINGS) NODE_CLASS_MAPPINGS.update(PURGE_VRAM_MAPPINGS) NODE_CLASS_MAPPINGS.update(CROP_MASK_MAPPINGS) NODE_CLASS_MAPPINGS.update(RESTORE_CROP_MAPPINGS) +NODE_CLASS_MAPPINGS.update(COLOR_MATCH_MAPPINGS) +NODE_CLASS_MAPPINGS.update(BBOX_VISUALIZE_MAPPINGS) NODE_CLASS_MAPPINGS.update(FILL_MASKED_MAPPINGS) -NODE_CLASS_MAPPINGS.update(MASK_OPERATIONS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(MASK_MAPPINGS) NODE_CLASS_MAPPINGS.update(LOAD_AUDIO_MAPPINGS) NODE_CLASS_MAPPINGS.update(AUDIO_CROP_MAPPINGS) NODE_CLASS_MAPPINGS.update(TEXTBOX_MAPPINGS) @@ -587,6 +628,7 @@ NODE_CLASS_MAPPINGS.update(THINK_REMOVER_MAPPINGS) NODE_CLASS_MAPPINGS.update(LORA_INFO_MAPPINGS) NODE_CLASS_MAPPINGS.update(KONTEXT_PRESETS_MAPPINGS) NODE_CLASS_MAPPINGS.update(PROMPT_HELPER_MAPPINGS) +NODE_CLASS_MAPPINGS.update(COLOR_TO_MASK_MAPPINGS) # 合并显示名称映射 NODE_DISPLAY_NAME_MAPPINGS = {} @@ -607,10 +649,12 @@ NODE_DISPLAY_NAME_MAPPINGS.update(CHECK_MASK_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(PURGE_VRAM_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(CROP_MASK_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(RESTORE_CROP_DISPLAY) +NODE_DISPLAY_NAME_MAPPINGS.update(COLOR_MATCH_DISPLAY) +NODE_DISPLAY_NAME_MAPPINGS.update(BBOX_VISUALIZE_DISPLAY) NODE_DISPLAY_NAME_MAPPINGS.update(FILL_MASKED_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_AUDIO_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(AUDIO_CROP_DISPLAY_MAPPINGS) -NODE_DISPLAY_NAME_MAPPINGS.update(MASK_OPERATIONS_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(MASK_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(TEXTBOX_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(TEXT_CONCATENATE_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(MATH_EXPRESSION_DISPLAY_MAPPINGS) @@ -618,6 +662,7 @@ NODE_DISPLAY_NAME_MAPPINGS.update(THINK_REMOVER_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(LORA_INFO_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(KONTEXT_PRESETS_DISPLAY_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(PROMPT_HELPER_DISPLAY_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(COLOR_TO_MASK_DISPLAY) NODE_CATEGORIES = { "UniversalToolkit": [ @@ -644,12 +689,15 @@ NODE_CATEGORIES = { "PurgeVRAM_UTK", "CropByMask_UTK", "RestoreCropBox_UTK", + "ColorMatch_UTK", + "BboxVisualize_UTK", "TextboxNode_UTK", "TextConcatenate_UTK", "MathExpression_UTK", "ThinkRemover_UTK", "LoraInfo_UTK", "LoadKontextPresets_UTK", + "ColorToMask_UTK", ] } diff --git a/nodes/image/bbox_visualize.py b/nodes/image/bbox_visualize.py new file mode 100644 index 0000000..eb47752 --- /dev/null +++ b/nodes/image/bbox_visualize.py @@ -0,0 +1,120 @@ +""" +Bbox Visualize Node for ComfyUI Universal Toolkit +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Bbox visualization functionality adapted from kjnodes. +Draws bounding boxes on images for visualization purposes. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import torch + +# Define BBOX type for ComfyUI +# BBOX is typically a list/tensor of [x, y, width, height] or [x1, y1, x2, y2] +# We'll register it as a custom type +if not hasattr(torch, 'BBOX'): + # Register BBOX as a custom type that can hold bbox coordinates + class BBOX: + pass + + +class BboxVisualize_UTK: + """ + Bbox Visualize node that draws bounding boxes on images. + + This node takes images and bounding box coordinates, then draws + rectangular frames on the images to visualize the bounding boxes. + Useful for debugging object detection, cropping operations, or + highlighting specific regions in images. + """ + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "visualizebbox" + CATEGORY = "UniversalToolkit/Image" + + DESCRIPTION = """ +Visualizes the specified bbox on the image. +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "bboxes": ("BBOX",), + "line_width": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), + "bbox_format": (["xywh", "xyxy"], {"default": "xywh"}), + }, + } + + def visualizebbox(self, bboxes, images, line_width, bbox_format): + """ + Visualizes the specified bbox on the image. + Adapted from kjnodes implementation. + """ + image_list = [] + for image, bbox in zip(images, bboxes): + if bbox_format == "xywh": + x_min, y_min, width, height = bbox + elif bbox_format == "xyxy": + x_min, y_min, x_max, y_max = bbox + width = x_max - x_min + height = y_max - y_min + else: + raise ValueError(f"Unknown bbox_format: {bbox_format}") + + # Ensure bbox coordinates are integers + x_min = int(x_min) + y_min = int(y_min) + width = int(width) + height = int(height) + + # Permute the image dimensions + image = image.permute(2, 0, 1) + + # Clone the image to draw bounding boxes + img_with_bbox = image.clone() + + # Define the color for the bbox, e.g., red + color = torch.tensor([1, 0, 0], dtype=torch.float32) + + # Ensure color tensor matches the image channels + if color.shape[0] != img_with_bbox.shape[0]: + color = color.unsqueeze(1).expand(-1, line_width) + + # Draw lines for each side of the bbox with the specified line width + for lw in range(line_width): + # Top horizontal line + if y_min + lw < img_with_bbox.shape[1]: + img_with_bbox[:, y_min + lw, x_min:x_min + width] = color[:, None] + + # Bottom horizontal line + if y_min + height - lw < img_with_bbox.shape[1]: + img_with_bbox[:, y_min + height - lw, x_min:x_min + width] = color[:, None] + + # Left vertical line + if x_min + lw < img_with_bbox.shape[2]: + img_with_bbox[:, y_min:y_min + height, x_min + lw] = color[:, None] + + # Right vertical line + if x_min + width - lw < img_with_bbox.shape[2]: + img_with_bbox[:, y_min:y_min + height, x_min + width - lw] = color[:, None] + + # Permute the image dimensions back + img_with_bbox = img_with_bbox.permute(1, 2, 0).unsqueeze(0) + image_list.append(img_with_bbox) + + return (torch.cat(image_list, dim=0),) + + +# Node registration +NODE_CLASS_MAPPINGS = { + "BboxVisualize_UTK": BboxVisualize_UTK, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "BboxVisualize_UTK": "Bbox Visualize (UTK)", +} diff --git a/nodes/image/color_match_standalone.py b/nodes/image/color_match_standalone.py new file mode 100644 index 0000000..fa53267 --- /dev/null +++ b/nodes/image/color_match_standalone.py @@ -0,0 +1,164 @@ +""" +Color Match Node for ComfyUI Universal Toolkit - Standalone Version +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Color matching functionality adapted from kjnodes. +This is a standalone version that follows the project's existing architecture. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import torch +import numpy as np +import os +from concurrent.futures import ThreadPoolExecutor + + +class ColorMatch_UTK: + """ + Color matching node that enables color transfer across images. + + This node is based on the color-matcher library and provides various methods + for transferring color characteristics from a reference image to a target image. + Useful for automatic color-grading of photographs, paintings, and film sequences. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_target": ("IMAGE", {"tooltip": "Target image to apply color matching to"}), + "image_ref": ("IMAGE", {"tooltip": "Reference image to match colors from"}), + "method": ( + [ + 'mkl', + 'hm', + 'reinhard', + 'mvgd', + 'hm-mvgd-hm', + 'hm-mkl-hm', + ], { + "default": 'mkl', + "tooltip": "Color matching method to use" + } + ), + }, + "optional": { + "strength": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 10.0, + "step": 0.01, + "tooltip": "Strength of the color matching effect" + }), + "multithread": ("BOOLEAN", { + "default": True, + "tooltip": "Use multithreading for batch processing" + }), + } + } + + CATEGORY = "UniversalToolkit/Image" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "colormatch" + + DESCRIPTION = """ +Color-matcher enables color transfer across images which comes in handy for automatic +color-grading of photographs, paintings and film sequences as well as light-field +and stopmotion corrections. + +The methods behind the mappings are based on the approach from Reinhard et al., +the Monge-Kantorovich Linearization (MKL) as proposed by Pitie et al. and our analytical solution +to a Multi-Variate Gaussian Distribution (MVGD) transfer in conjunction with classical histogram +matching. As shown below our HM-MVGD-HM compound outperforms existing methods. + +Methods: +- mkl: Monge-Kantorovich Linearization +- hm: Histogram Matching +- reinhard: Reinhard et al. method +- mvgd: Multi-Variate Gaussian Distribution +- hm-mvgd-hm: Histogram Matching + MVGD + Histogram Matching +- hm-mkl-hm: Histogram Matching + MKL + Histogram Matching + +Reference: https://github.com/hahnec/color-matcher/ +""" + + def colormatch(self, image_target, image_ref, method, strength=1.0, multithread=True): + """ + Apply color matching from reference image to target image. + + Args: + image_target: Target image tensor to be color matched + image_ref: Reference image tensor + method: Color matching method to use + strength: Strength of the color matching effect (0.0 to 10.0) + multithread: Whether to use multithreading for batch processing + + Returns: + Tuple containing the color matched image tensor + """ + try: + from color_matcher import ColorMatcher + except ImportError: + raise Exception( + "Can't import color-matcher. Please install it using: pip install color-matcher" + ) + + # Move tensors to CPU for processing + image_ref = image_ref.cpu() + image_target = image_target.cpu() + batch_size = image_target.size(0) + + # Remove batch dimension if single image + images_target = image_target.squeeze() + images_ref = image_ref.squeeze() + + # Convert to numpy arrays + image_ref_np = images_ref.numpy() + images_target_np = images_target.numpy() + + def process(i): + """Process a single image in the batch.""" + cm = ColorMatcher() + + # Handle batch vs single image + image_target_np_i = images_target_np if batch_size == 1 else images_target[i].numpy() + image_ref_np_i = image_ref_np if image_ref.size(0) == 1 else images_ref[i].numpy() + + try: + # Apply color matching + image_result = cm.transfer(src=image_target_np_i, ref=image_ref_np_i, method=method) + + # Apply strength blending + image_result = image_target_np_i + strength * (image_result - image_target_np_i) + + return torch.from_numpy(image_result) + except Exception as e: + print(f"Color matching error for image {i}: {e}") + return torch.from_numpy(image_target_np_i) # Return original as fallback + + # Process images (with or without multithreading) + if multithread and batch_size > 1: + max_threads = min(os.cpu_count() or 1, batch_size) + with ThreadPoolExecutor(max_workers=max_threads) as executor: + out = list(executor.map(process, range(batch_size))) + else: + out = [process(i) for i in range(batch_size)] + + # Stack results and ensure proper format + out = torch.stack(out, dim=0).to(torch.float32) + out.clamp_(0, 1) # Ensure values are in valid range + + return (out,) + + +# Node registration - Following the project's existing pattern +NODE_CLASS_MAPPINGS = { + "ColorMatch_UTK": ColorMatch_UTK, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ColorMatch_UTK": "Color Match (UTK)", +} diff --git a/nodes/mask/__init__.py b/nodes/mask/__init__.py new file mode 100644 index 0000000..efeb581 --- /dev/null +++ b/nodes/mask/__init__.py @@ -0,0 +1,29 @@ +""" +ComfyUI Universal Toolkit - Mask Nodes +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Mask processing nodes for ComfyUI Universal Toolkit. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import importlib +import os + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# 自动导入本目录下所有节点文件的注册表 +for filename in os.listdir(os.path.dirname(__file__)): + if filename.endswith(".py") and filename not in ( + "__init__.py", + ): + modulename = filename[:-3] + module = importlib.import_module(f".{modulename}", __package__) + if hasattr(module, "NODE_CLASS_MAPPINGS"): + NODE_CLASS_MAPPINGS.update(getattr(module, "NODE_CLASS_MAPPINGS")) + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS.update( + getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") + ) diff --git a/nodes/mask/separate_masks.py b/nodes/mask/separate_masks.py new file mode 100644 index 0000000..e8560e8 --- /dev/null +++ b/nodes/mask/separate_masks.py @@ -0,0 +1,237 @@ +""" +Separate Masks Node for ComfyUI Universal Toolkit +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Separate masks functionality adapted from kjnodes. +Separates a mask into multiple masks based on connected components. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import torch +import numpy as np +from comfy.utils import ProgressBar + + +class SeparateMasks_UTK: + """ + Separate Masks node that divides a mask into multiple masks based on connected components. + + This node analyzes input masks and separates them into individual masks for each + connected component that meets the size threshold requirements. Useful for isolating + different objects or regions within a single mask. + """ + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("masks",) + FUNCTION = "separate_masks" + CATEGORY = "UniversalToolkit/Mask" + OUTPUT_NODE = True + + DESCRIPTION = """ +Separates a mask into multiple masks based on connected components. + +The node identifies connected components (continuous regions) in the input mask +and creates separate masks for each component that meets the size thresholds. +Components are sorted by their horizontal position (left to right). + +Modes: +- **area**: Preserves the exact shape of each component +- **box**: Creates rectangular bounding boxes around each component +- **convex_polygons**: Creates simplified convex polygon approximations + +Size thresholds filter out small noise or unwanted components. +Components smaller than the specified width/height are ignored. + +Useful for: +- Separating multiple objects in a single mask +- Filtering components by size +- Creating individual masks for batch processing +- Object isolation and analysis +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK", {"tooltip": "Input mask to separate into components"}), + "size_threshold_width": ("INT", { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "Minimum width for components to be included" + }), + "size_threshold_height": ("INT", { + "default": 256, + "min": 0, + "max": 4096, + "step": 1, + "tooltip": "Minimum height for components to be included" + }), + "mode": (["area", "box", "convex_polygons"], { + "default": "area", + "tooltip": "Method for creating separated masks" + }), + "max_poly_points": ("INT", { + "default": 8, + "min": 3, + "max": 32, + "step": 1, + "tooltip": "Maximum points for polygon approximation (convex_polygons mode)" + }), + }, + } + + def polygon_to_mask(self, polygon, shape): + """Convert polygon points to mask.""" + try: + import cv2 + except ImportError: + raise Exception("OpenCV is required for polygon operations. Please install: pip install opencv-python") + + mask = np.zeros((shape[0], shape[1]), dtype=np.uint8) + + if len(polygon.shape) == 2: # Check if polygon points are valid + polygon = polygon.astype(np.int32) + cv2.fillPoly(mask, [polygon], 1) + return mask + + def get_mask_polygon(self, mask_np, max_points): + """Extract polygon approximation from mask.""" + try: + import cv2 + except ImportError: + raise Exception("OpenCV is required for polygon operations. Please install: pip install opencv-python") + + contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + return None + + largest_contour = max(contours, key=cv2.contourArea) + hull = cv2.convexHull(largest_contour) + + # Initialize with smaller epsilon for more points + perimeter = cv2.arcLength(hull, True) + epsilon = perimeter * 0.01 # Start smaller + + min_eps = perimeter * 0.001 # Much smaller minimum + max_eps = perimeter * 0.2 # Smaller maximum + + best_approx = None + best_diff = float('inf') + max_iterations = 20 + + for i in range(max_iterations): + curr_eps = (min_eps + max_eps) / 2 + approx = cv2.approxPolyDP(hull, curr_eps, True) + points_diff = len(approx) - max_points + + if abs(points_diff) < best_diff: + best_approx = approx + best_diff = abs(points_diff) + + if len(approx) > max_points: + min_eps = curr_eps * 1.1 # More gradual adjustment + elif len(approx) < max_points: + max_eps = curr_eps * 0.9 # More gradual adjustment + else: + return approx.squeeze() + + if abs(max_eps - min_eps) < perimeter * 0.0001: # Relative tolerance + break + + # If we didn't find exact match, return best approximation + return best_approx.squeeze() if best_approx is not None else hull.squeeze() + + def separate_masks(self, mask, size_threshold_width, size_threshold_height, mode, max_poly_points): + """ + Separate mask into individual component masks. + + Args: + mask: Input mask tensor + size_threshold_width: Minimum width for components + size_threshold_height: Minimum height for components + mode: Separation mode ('area', 'box', 'convex_polygons') + max_poly_points: Maximum points for polygon approximation + + Returns: + Tuple containing separated masks tensor + """ + try: + from scipy.ndimage import label + except ImportError: + raise Exception("SciPy is required for connected component analysis. Please install: pip install scipy") + + B, H, W = mask.shape + separated = [] + + mask = mask.round() + + for b in range(B): + mask_np = mask[b].cpu().numpy().astype(np.uint8) + structure = np.ones((3, 3), dtype=np.int8) + labeled, ncomponents = label(mask_np, structure=structure) + pbar = ProgressBar(ncomponents) + + for component in range(1, ncomponents + 1): + component_mask_np = (labeled == component).astype(np.uint8) + + # Find bounding box + rows = np.any(component_mask_np, axis=1) + cols = np.any(component_mask_np, axis=0) + + if not np.any(rows) or not np.any(cols): + pbar.update(1) + continue + + y_min, y_max = np.where(rows)[0][[0, -1]] + x_min, x_max = np.where(cols)[0][[0, -1]] + + width = x_max - x_min + 1 + height = y_max - y_min + 1 + centroid_x = (x_min + x_max) / 2 # Calculate x centroid + + print(f"Component {component}: width={width}, height={height}, x_pos={centroid_x}") + + # Check size thresholds + if width >= size_threshold_width and height >= size_threshold_height: + if mode == "convex_polygons": + polygon = self.get_mask_polygon(component_mask_np, max_poly_points) + if polygon is not None: + poly_mask = self.polygon_to_mask(polygon, (H, W)) + poly_mask = torch.tensor(poly_mask, device=mask.device, dtype=torch.float32) + separated.append((centroid_x, poly_mask)) + elif mode == "box": + # Create bounding box mask + box_mask = np.zeros((H, W), dtype=np.uint8) + box_mask[y_min:y_max+1, x_min:x_max+1] = 1 + box_mask = torch.tensor(box_mask, device=mask.device, dtype=torch.float32) + separated.append((centroid_x, box_mask)) + else: # mode == "area" + area_mask = torch.tensor(component_mask_np, device=mask.device, dtype=torch.float32) + separated.append((centroid_x, area_mask)) + + pbar.update(1) + + if len(separated) > 0: + # Sort by x position and extract only the masks + separated.sort(key=lambda x: x[0]) + separated = [x[1] for x in separated] + out_masks = torch.stack(separated, dim=0) + return (out_masks,) + else: + # Return empty mask if no components found + empty_mask = torch.zeros((1, H, W), device=mask.device, dtype=torch.float32) + return (empty_mask,) + + +# Node registration +NODE_CLASS_MAPPINGS = { + "SeparateMasks_UTK": SeparateMasks_UTK, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SeparateMasks_UTK": "Separate Masks (UTK)", +} diff --git a/nodes/tools/color_to_mask.py b/nodes/tools/color_to_mask.py new file mode 100644 index 0000000..df0a231 --- /dev/null +++ b/nodes/tools/color_to_mask.py @@ -0,0 +1,164 @@ +""" +Color to Mask Node for ComfyUI Universal Toolkit +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Color to mask conversion functionality adapted from kjnodes. +Converts chosen RGB values to mask based on color distance threshold. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import torch +from comfy.utils import ProgressBar + + +class ColorToMask_UTK: + """ + Color to Mask node that converts chosen RGB values to mask. + + This node analyzes input images and creates masks based on color similarity + to a specified RGB color within a threshold distance. Useful for isolating + specific colored areas in images for masking purposes. + """ + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "color_to_mask" + CATEGORY = "UniversalToolkit/Tools" + + DESCRIPTION = """ +Converts chosen RGB value to a mask based on color distance. + +The node calculates the Euclidean distance between each pixel's color +and the target RGB color. Pixels within the threshold distance are +included in the mask (white), while others are excluded (black). + +With batch inputs, the **per_batch** parameter controls the number +of images processed at once for memory efficiency. + +Parameters: +- RGB values: Target color to match (0-255) +- Threshold: Maximum color distance for inclusion (0-255) +- Invert: Flip the mask (exclude matching colors instead) +- Per batch: Number of images to process simultaneously +""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", {"tooltip": "Input images to convert to masks"}), + "red": ("INT", { + "default": 0, + "min": 0, + "max": 255, + "step": 1, + "tooltip": "Red component of target color" + }), + "green": ("INT", { + "default": 0, + "min": 0, + "max": 255, + "step": 1, + "tooltip": "Green component of target color" + }), + "blue": ("INT", { + "default": 0, + "min": 0, + "max": 255, + "step": 1, + "tooltip": "Blue component of target color" + }), + "threshold": ("INT", { + "default": 10, + "min": 0, + "max": 255, + "step": 1, + "tooltip": "Maximum color distance for inclusion in mask" + }), + "invert": ("BOOLEAN", { + "default": False, + "tooltip": "Invert the mask (exclude matching colors)" + }), + "per_batch": ("INT", { + "default": 16, + "min": 1, + "max": 4096, + "step": 1, + "tooltip": "Number of images to process simultaneously" + }), + }, + } + + def color_to_mask(self, images, red, green, blue, threshold, invert, per_batch): + """ + Convert RGB color to mask based on color distance threshold. + + Args: + images: Input image tensor + red: Red component of target color (0-255) + green: Green component of target color (0-255) + blue: Blue component of target color (0-255) + threshold: Maximum color distance for inclusion + invert: Whether to invert the mask + per_batch: Number of images to process per batch + + Returns: + Tuple containing the mask tensor + """ + # Define target color and mask colors + 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) + + # Swap colors if invert is True + if invert: + black, white = white, black + + # Initialize progress bar and output list + steps = images.shape[0] + pbar = ProgressBar(steps) + tensors_out = [] + + # Process images in batches + for start_idx in range(0, images.shape[0], per_batch): + end_idx = min(start_idx + per_batch, images.shape[0]) + batch_images = images[start_idx:end_idx] + + # Calculate color distances using Euclidean distance + # Convert images from [0,1] to [0,255] range for comparison + color_distances = torch.norm(batch_images * 255 - color.float(), dim=-1) + + # Create mask based on threshold + mask = color_distances <= threshold + + # Apply mask to create output (white for match, black for no match) + mask_out = torch.where(mask.unsqueeze(-1), white.float(), black.float()) + + # Convert to grayscale mask by taking mean across color channels + mask_out = mask_out.mean(dim=-1) + + # Normalize to [0,1] range and move to CPU + mask_out = mask_out / 255.0 + tensors_out.append(mask_out.cpu()) + + # Update progress bar + batch_count = mask_out.shape[0] + pbar.update(batch_count) + + # Concatenate all batches and ensure proper range + tensors_out = torch.cat(tensors_out, dim=0) + tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0) + + return (tensors_out,) + + +# Node registration +NODE_CLASS_MAPPINGS = { + "ColorToMask_UTK": ColorToMask_UTK, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ColorToMask_UTK": "Color To Mask (UTK)", +} diff --git a/nodes/tools/kontext_presets.py b/nodes/tools/kontext_presets.py index 86167cb..3f2a31b 100644 --- a/nodes/tools/kontext_presets.py +++ b/nodes/tools/kontext_presets.py @@ -80,7 +80,7 @@ class LoadKontextPresets_UTK: }, { "name": "Ecommerce-Fashion Try-On Model Showcase (时尚试穿展示)", - "brief": "Transform a clothing item image into a realistic and stylish model showcase scene based on $user_prompt$. Generate a professional model wearing or presenting the clothing item in a natural pose within an appropriate lifestyle or fashion setting. The clothing item's structure, color, and design must remain completely unchanged. Apply the following principles: 1. **Clothing-to-Body Mapping**: Apply the clothing image to a human model in a way that naturally fits the body, respecting the shape, folds, and texture of the original item. Avoid distortion or unrealistic deformation while ensuring proper fit and drape. 2. **Pose & Fit Realism**: Select a natural, stylish pose that best showcases the clothing's cut and design. Ensure correct tension points, fabric drape, and body posture based on the item type (e.g., jacket, dress, hoodie, formal wear). 3. **Contextual Scene Adaptation**: Generate a lifestyle or fashion-oriented background that matches $user_prompt$ (e.g., urban streetwear, elegant indoors, casual beach, fashion runway, studio setting). Ensure the background supports but doesn't overpower the clothing item. 4. **Professional Lighting & Composition**: Use natural or studio lighting based on the context. Ensure shadows, highlights, and color tones are realistic and enhance the garment's appeal while maintaining the original clothing colors and textures. 5. **Model Presentation**: Ensure the model has human-proportioned body with stylish appearance and varied ethnicity options. Avoid excessive stylization unless explicitly requested in $user_prompt$. 6. **E-commerce Quality**: Create images suitable for e-commerce, lookbook, fashion campaign, or editorial use depending on the scene input. Maintain commercial-grade quality with sharp focus on the clothing item. 7. **Flexible Customization**: Allow background, model pose, and styling to be varied via $user_prompt$ for different use cases while preserving the clothing item's integrity.", + "brief": "Transform a clothing item image into a realistic and stylish model showcase scene based on $user_prompt$. Generate a professional model wearing or presenting the clothing item in a natural pose within an appropriate lifestyle or fashion setting. The clothing item's structure, color, and design must remain completely unchanged. Apply the following principles: 1. **Clothing-to-Body Mapping**: Apply the clothing image to a human model in a way that naturally fits the body, respecting the shape, folds, and texture of the original item. Avoid distortion or unrealistic deformation while ensuring proper fit and drape. 2. **Model Customization via Prompt**: Use $user_prompt$ to specify model characteristics such as age, gender, ethnicity, body type, hairstyle, and overall appearance. Support diverse model representation including different ages (young adult, mature, senior), genders (male, female, non-binary), ethnicities, and body types (athletic, curvy, slim, plus-size). 3. **Scene & Environment Control**: Generate backgrounds and settings based on $user_prompt$ specifications. Support various environments including urban streetwear (city streets, graffiti walls), elegant indoors (luxury hotel, modern apartment), casual beach (seaside, beachfront), fashion runway (catwalk, studio), outdoor lifestyle (park, café, rooftop), and seasonal settings (spring garden, winter cityscape). 4. **Pose & Styling Flexibility**: Allow pose and styling variations through $user_prompt$. Support different poses (standing, walking, sitting, dynamic movement), camera angles (front view, side profile, three-quarter view), and styling elements (accessories, makeup, hair styling) that complement the clothing item. 5. **Professional Lighting & Composition**: Use natural or studio lighting based on the context specified in $user_prompt$. Ensure shadows, highlights, and color tones are realistic and enhance the garment's appeal while maintaining the original clothing colors and textures. 6. **E-commerce Quality**: Create images suitable for e-commerce, lookbook, fashion campaign, or editorial use depending on the scene input. Maintain commercial-grade quality with sharp focus on the clothing item. 7. **Dynamic Adaptation**: Fully utilize $user_prompt$ to customize all aspects including model appearance, scene setting, pose selection, lighting style, and overall mood while preserving the clothing item's integrity and design.", }, { "name": "Ecommerce-Product Pattern Extraction (产品图案提取)", @@ -88,7 +88,7 @@ class LoadKontextPresets_UTK: }, { "name": "Ecommerce-Logo Transfer to Product (品牌融合植入)", - "brief": "Seamlessly blend a logo from one image onto a product in another image based on $user_prompt$, ensuring that the product's structure, texture, lighting, and composition remain completely unchanged. The logo must appear naturally integrated, as if originally part of the product. Apply the following principles: 1. **Logo Extraction & Preservation**: Isolate the logo from the source image with high fidelity, preserving its shape, color, proportions, and any visual branding characteristics while maintaining its original design integrity. 2. **Product Integrity Protection**: Do not alter the original product photo in any way — its composition, lighting, texture, material, colors, shadows, and surface details must remain exactly as they appear in the target image. 3. **Realistic Logo Fusion**: Integrate the logo into the product image in a way that fully respects the surface curvature, fabric or material behavior (e.g., cloth folds, leather texture, matte or glossy reflection). Avoid floating, tiling, or artificial appearance. 4. **Lighting & Perspective Matching**: The logo must match the product's lighting and perspective perfectly. If the product is lit from the top-left, the logo must reflect that same direction of light, shadow softness, and reflection if applicable. 5. **Precise Placement Control**: Use $user_prompt$ to determine the intended placement and scale of the logo (e.g., 'top-right corner of the backpack flap', 'center chest of the t-shirt', or 'front face of the bottle'). 6. **Non-Obstructive Integration**: Avoid covering essential product details, existing logos, or stitching unless explicitly stated. The logo should be placed where it naturally fits into the design. 7. **Commercial Presentation Standards**: The final image should appear like a real professional product photo prepared for use in brand catalogs, ecommerce platforms, or promotional displays — high-resolution, polished, and photorealistic. NOTE: This preset requires two input images - one containing the logo to be transferred, and one containing the target product.", + "brief": "Transfer a logo from one image onto a product in another image based on $user_prompt$. Ensure the product image remains completely unchanged, including all details, structure, textures, lighting, and background. The logo must be seamlessly integrated as if it were originally part of the product. Apply the following principles: 1. **Preserve Product and Background**: Do not alter any aspect of the product photo, including the product itself and its background. Maintain exact composition, lighting, textures, shadows, materials, and any scene elements. No hallucination, repainting, or cleanup is allowed. 2. **Logo Extraction and Fidelity**: Extract the logo from the source image precisely, retaining its shape, proportions, colors, and visual style. Remove any surrounding background if needed without distorting the logo. 3. **Seamless Logo Integration**: Integrate the logo onto the product surface in a photorealistic manner. The logo must conform to the product's surface (e.g., wrapping on fabric, curving over a bottle), with accurate distortion, reflection, and occlusion if applicable. 4. **Lighting and Perspective Match**: The logo must match the product image's lighting direction, intensity, softness, and perspective. It must appear as part of the original lighting environment. 5. **Precise Placement Control**: Follow the user-defined placement in $user_prompt$. Example: 'top-left of the backpack flap', 'center of the shirt chest', 'side of the shoebox', etc. 6. **No Visual Conflicts**: Avoid placing the logo over existing branding, design elements, stitching, or essential product details — unless the prompt allows it. 7. **Commercial Realism**: The final image must meet the quality standards of professional product photography. The result should appear as a genuine product photo with the brand logo naturally embedded in the design, suitable for e-commerce or advertising. NOTE: This preset requires two input images - one containing the logo to be transferred, and one containing the target product.", }, # === 人物变换类 === { diff --git a/pyproject.toml b/pyproject.toml index 32ef18e..7cc7824 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "universaltoolkit" description = "A comprehensive toolkit based on ComfyUI, providing image, mask, audio, and tools nodes, fully modular and v3 compatible." -version = "1.3.2" +version = "1.3.3" license = {file = "LICENSE"} dependencies = [ "torch", diff --git a/requirements.txt b/requirements.txt index c226b0f..afc1c0c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,6 +12,9 @@ soundfile>=0.10.0 # Scientific computing scipy>=1.7.0 +# Color matching +color-matcher + # Network requests requests>=2.25.0 aiohttp>=3.8.0