diff --git a/README.MD b/README.MD
index 03a58f1..ad2674e 100644
--- a/README.MD
+++ b/README.MD
@@ -70,6 +70,7 @@ When this error has occurred, please check the network environment.
## Update
**If the dependency package error after updating, please reinstall the relevant dependency packages.
+* Commit [MaskByColor](#MaskByColor) node, Generate a mask based on the selected color.
* Commit [LoadPSD](#LoadPSD) node, It read the psd format, and output layer images.
* Commit [SegformerB2ClothesUltra](#SegformerB2ClothesUltra) node, it used to segment character clothing. The model segmentation code is from[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes), thanks to the original author.
* [SaveImagePlus](#SaveImagePlus) node adds the output workflow to the json function, supports ```%date``` and ```%time``` to embeddint date or time to path and filename, and adds the preview switch.
@@ -1374,6 +1375,20 @@ Outputs:
* yolo_masks: For all masks identified by yolo, each individual mask is output as a mask.
+### MaskByColor
+Generate a mask based on the selected color.
+
+
+Node Options:
+
+* image: Input image.
+* color: Color selector. Click on the color block to select a color, and you can use the straws on the color picker panel to pick up the screen color. Note: When using straws, maximize the browser window.
+* color_in_HEX4: Enter color values. If this item has input, it will be used first, ignoring the color selected by the ```color``` .
+* threshold: Mask range threshold, the larger the value, the larger the mask range.
+* fix_gap: Repair the gaps in the mask. If there are obvious gaps in the mask, increase this value appropriately.
+* fix_threshold: The threshold for repairing masks.
+* invert_mask: Whether to reverse the mask.
+
### Shadow & Highlight Mask
Generate masks for the dark and bright parts of the image.
diff --git a/README_CN.MD b/README_CN.MD
index ce3e55c..789ad52 100644
--- a/README_CN.MD
+++ b/README_CN.MD
@@ -70,6 +70,7 @@ git clone https://github.com/chflame163/ComfyUI_LayerStyle.git
## 更新说明
**如果本插件更新后出现依赖包错误，请重新安装相关依赖包。
+* 添加 [MaskByColor](#MaskByColor) 节点, 根据选择的颜色生成遮罩。
* 添加 [LoadPSD](#LoadPSD) 节点, 读取psd格式并输出图层图片。
* 添加 [SegformerB2ClothesUltra](#SegformerB2ClothesUltra)节点，用于分割人物服装。模型分割代码来自[StartHua](https://github.com/StartHua/Comfyui_segformer_b2_clothes)，感谢原作者。
* [SaveImagePlus](#SaveImagePlus)节点增加输出工作流为json功能，支持使用```%date```和```%time```在路径和文件名嵌入时间，增加预览开关。
@@ -1360,6 +1361,21 @@ PersonMaskUltra的V2升级版，增加了VITMatte边缘处理方法。(注意:
* yolo_masks: yolo识别出来的所有遮罩，每个单独的遮罩输出为一个mask。
+### MaskByColor
+根据颜色生成遮罩。
+
+
+节点选项说明:
+
+* image: 图像输入。
+* color: 颜色选择器。点击色块选择颜色，可以使用选色器面板上的吸管拾取屏幕颜色。注意:使用吸管时，需将浏览器窗口最大化。
+* color_in_HEX4: 输入色值。如果此项有输入，则优先使用，忽略```color```选取的颜色。
+* threshold: 遮罩范围阈值，数值越大，遮罩范围越大。
+* fix_gap: 修补遮罩中的空隙。如果遮罩中有比较明显的空隙，适当调高此数值。
+* fix_threshold: 修补遮罩的阈值。
+* invert_mask: 是否反转遮罩。
+
+
### Shadow & Highlight Mask
生成图像暗部和亮部的遮罩。
diff --git a/image/mask_by_color_example.jpg b/image/mask_by_color_example.jpg
new file mode 100644
index 0000000..4dcb74c
Binary files /dev/null and b/image/mask_by_color_example.jpg differ
diff --git a/image/mask_by_color_node.jpg b/image/mask_by_color_node.jpg
new file mode 100644
index 0000000..7a0e341
Binary files /dev/null and b/image/mask_by_color_node.jpg differ
diff --git a/py/imagefunc.py b/py/imagefunc.py
index 38c48f1..d5dbb41 100644
--- a/py/imagefunc.py
+++ b/py/imagefunc.py
@@ -1262,6 +1262,40 @@ def generate_text_image(width:int, height:int, text:str, font_file:str, text_sca
return image
'''Mask Functions'''
+
+def create_mask_from_color_cv2(image:Image, color:str, tolerance:int=0) -> Image:
+ (r, g, b) = Hex_to_RGB(color)
+ target_color = (b, g, r)
+ tolerance = 127 + int(tolerance * 1.28)
+ # tolerance = 255 - tolerance
+ # 将RGB颜色转换为HSV颜色空间
+ image = pil2cv2(image)
+ hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+
+ # 定义目标颜色的HSV范围
+ lower_color = np.array([max(target_color[0] - tolerance, 0), max(target_color[1] - tolerance, 0), max(target_color[2] - tolerance, 0)])
+ upper_color = np.array([min(target_color[0] + tolerance, 255), min(target_color[1] + tolerance, 255), min(target_color[2] + tolerance, 255)])
+
+ # 创建掩码
+ mask = cv2.inRange(hsv_image, lower_color, upper_color)
+
+ return cv22pil(mask).convert("L")
+
+def create_mask_from_color_tensor(image:Image, color:str, tolerance:int=0) -> Image:
+ threshold = int(tolerance * 1.28)
+ (red, green, blue) = Hex_to_RGB(color)
+ image = pil2tensor(image).squeeze()
+ temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)
+ color_value = torch.tensor([red, green, blue])
+ lower_bound = (color_value - threshold).clamp(min=0)
+ upper_bound = (color_value + threshold).clamp(max=255)
+ lower_bound = lower_bound.view(1, 1, 1, 3)
+ upper_bound = upper_bound.view(1, 1, 1, 3)
+ mask = (temp >= lower_bound) & (temp <= upper_bound)
+ mask = mask.all(dim=-1)
+ mask = mask.float()
+ return tensor2pil(mask).convert("L")
+
@lru_cache(maxsize=1, typed=False)
def load_RMBG_model():
current_directory = os.path.dirname(os.path.abspath(__file__))
diff --git a/py/mask_by_color.py b/py/mask_by_color.py
new file mode 100644
index 0000000..03e253d
--- /dev/null
+++ b/py/mask_by_color.py
@@ -0,0 +1,54 @@
+from .imagefunc import *
+
+NODE_NAME = 'MaskByColor'
+
+
+class MaskByColor:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {
+ "required": {
+ "image": ("IMAGE", ),
+ "color": ("COLOR", {"default": "#FFFFFF"},),
+ "color_in_HEX": ("STRING", {"default": ""}),
+ "threshold": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1, }),
+ "fix_gap": ("INT", {"default": 2, "min": 0, "max": 32, "step": 1}),
+ "fix_threshold": ("FLOAT", {"default": 0.75, "min": 0.01, "max": 0.99, "step": 0.01}),
+ "invert_mask": ("BOOLEAN", {"default": False}), # 反转mask
+ }
+ }
+
+ RETURN_TYPES = ("MASK",)
+ RETURN_NAMES = ("mask",)
+ FUNCTION = "mask_by_color"
+ CATEGORY = '😺dzNodes/LayerMask'
+
+ def mask_by_color(self, image, color, color_in_HEX, threshold,
+ fix_gap, fix_threshold, invert_mask):
+
+ if color_in_HEX != "" and color_in_HEX.startswith('#') and len(color_in_HEX) == 7:
+ color = color_in_HEX
+
+ ret_masks = []
+
+ for i in image:
+ img = tensor2pil(i.unsqueeze(0))
+ mask = create_mask_from_color_tensor(img, color, threshold)
+ mask = image2mask(mask)
+ if invert_mask:
+ mask = 1 - mask
+
+ if fix_gap:
+ mask = mask_fix(mask, 1, fix_gap, fix_threshold, fix_threshold)
+ ret_masks.append(mask)
+
+ return (torch.cat(ret_masks, dim=0), )
+
+
+NODE_CLASS_MAPPINGS = {
+ "LayerMask: MaskByColor": MaskByColor
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "LayerMask: MaskByColor": "LayerMask: Mask by Color"
+}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 4436be5..48258ef 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
[project]
name = "comfyui_layerstyle"
description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress."
-version = "1.0.1"
+version = "1.0.2"
license = "LICENSE"
dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "opencv-contrib-python", "pymatting", "segment_anything", "timm", "addict", "yapf", "colour-science", "wget", "mediapipe", "loguru", "typer_config", "fastapi", "rich", "google-generativeai", "diffusers", "omegaconf", "tqdm", "transformers", "kornia", "image-reward", "ultralytics", "blend_modes", "blind-watermark", "qrcode", "pyzbar", "psd-tools"]