From d8b633ae61d1ae72dd2a00118c20917d43ae756c Mon Sep 17 00:00:00 2001 From: chflame163 <130118553+chflame163@users.noreply.github.com> Date: Tue, 16 Jan 2024 18:18:26 +0800 Subject: [PATCH] Word Cloud node add mask output --- py/comfy_wordcloud.py | 41 +++++++++++++++++++---------------------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/py/comfy_wordcloud.py b/py/comfy_wordcloud.py index 9d30973..4be10c8 100644 --- a/py/comfy_wordcloud.py +++ b/py/comfy_wordcloud.py @@ -1,3 +1,4 @@ +import math import os import glob import re @@ -5,14 +6,14 @@ import numpy as np import torch import matplotlib.pyplot as plt from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator -from PIL import Image, ImageOps +from PIL import Image, ImageChops import jieba def log(message): name = 'WordCloud' print(f"# 😺dzNodes: {name} -> {message}") -COLOR_MAP = ['viridis', 'Accent', 'Blues', 'BrBG', 'BuGn', 'BuPu', 'CMRmap','Dark2', 'GnBu', +COLOR_MAP = ['viridis', 'Accent', 'Blues', 'BrBG', 'BuGn', 'BuPu', 'CMRmap', 'Dark2', 'GnBu', 'Grays', 'Greens', 'OrRd', 'Oranges', 'PRGn', 'Paired', 'Pastel1', 'Pastel2', 'PiYG', 'PuBu', 'PuBuGn', 'PuOr', 'PuRd', 'Purples', 'RdBu', 'RdGy', 'RdPu', 'RdYlBu', 'RdYlGn', 'Reds', 'Set1', 'Set2', 'Set3', 'Spectral', 'Wistia', @@ -32,9 +33,9 @@ ini_file = os.path.join(os.path.dirname(os.path.dirname(os.path.normpath(__file_ try: with open(ini_file, 'r') as f: ini = f.read() - dir = ini[ini.find('=') + 1:].rstrip().lstrip() - if os.path.exists(dir): - font_dir = dir + d = ini[ini.find('=') + 1:].rstrip().lstrip() + if os.path.exists(d): + font_dir = d else: log(f'ERROR: invalid dir, default to be used. check {ini_file}') except Exception as e: @@ -62,8 +63,8 @@ def img_whitebackground(image): image = image.convert('RGBA') width = image.width height = image.height - img_new = Image.new('RGB',size=(width,height),color=(255,255,255)) - img_new.paste(image,(0,0),mask=image) + img_new = Image.new('RGB', size=(width, height), color=(255, 255, 255)) + img_new.paste(image, (0, 0), mask=image) return img_new @@ -83,30 +84,22 @@ class ComfyWordCloud: "height": ("INT", {"default": 512}), # 画幅高 "scale": ("FLOAT", {"default": 1, "min": 0.1, "max": 1000.0, "step": 0.01}), # 放大倍数 "margin": ("INT", {"default": 0}), # 空白边界 - ## font "font_path": (font_list,), # 字体文件 "min_font_size": ("INT", {"default": 4}), # 单词最小size "max_font_size": ("INT", {"default": 128}), # 单词最大size - # "font_step": ("INT", {"default": 1}), # 字体迭代步长，大于1时计算速度加快但易导致错误 "relative_scaling": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), # 单词大小离散度 ## color control "colormap": (COLOR_MAP,), # 文字颜色 "background_color": ("STRING", {"default": "#FFFFFF"}), # 背景颜色 - "transparent_background": ("BOOLEAN", {"default": False}), # 是否透明，如果是则需要background_color强制为None + "transparent_background": ("BOOLEAN", {"default": True}), # 是否透明，如果是则需要background_color强制为None ## word control "prefer_horizontal": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01}), # 横排比例 "max_words": ("INT", {"default": 200}), # 最大单词数量 "repeat": ("BOOLEAN", {"default": False}), # 允许重复单词直到最大单词数量 - # "min_word_length": ("INT", {"default": 0}), # 最小单词长度 "include_numbers": ("BOOLEAN", {"default": False}), # 是否包含数字 - # "collocations": ("BOOLEAN", {"default": False}), # 词组关联开关 - # "collocation_threshold": ("INT", {"default": 30}), # 词组关联度 - # "normalize_plurals": ("BOOLEAN", {"default": True}), # 复数单词转单数 - # "ranks_only": ("BOOLEAN", {"default": False}), # 仅显示高频词 "random_state": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}), # 固定随机值，-1时强制转为None（随机） "stopwords": ("STRING", {"default": ""}), # 排除词，用中英文逗号或空格分开 - # "regexp": ("STRING", {"default": "", "multiline": True}), # 正则表达式 string or None }, "optional": { ## recolor refrence image @@ -120,8 +113,8 @@ class ComfyWordCloud: } } - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) + RETURN_TYPES = ("IMAGE", "MASK",) + RETURN_NAMES = ("image", "mask",) FUNCTION = 'wordcloud' CATEGORY = '😺dzNodes' OUTPUT_NODE = True @@ -161,16 +154,17 @@ class ComfyWordCloud: stopwords_set = set("") if not stopwords == "": stopwords_list = re.split(r'[，,\s*]', stopwords) - stopwords_set= set([x for x in stopwords_list if x != '']) # 去除空字符 + stopwords_set = set([x for x in stopwords_list if x != '']) # 去除空字符 # 同时在词典中删除（stopwords之bug） for item in stopwords_set: if item in freq_dict.keys(): del freq_dict[item] + bg_color = background_color mode = 'RGB' if transparent_background: - background_color = None + bg_color = None mode = 'RGBA' if random_state == -1: @@ -190,7 +184,7 @@ class ComfyWordCloud: wc = WordCloud(width=width, height=height, scale=scale, margin=margin, font_path=font_path, min_font_size=min_font_size, max_font_size=max_font_size, relative_scaling=relative_scaling, colormap=colormap, mode=mode, - background_color=background_color, prefer_horizontal=prefer_horizontal, + background_color=bg_color, prefer_horizontal=prefer_horizontal, max_words=max_words, repeat=repeat, include_numbers=include_numbers, random_state=random_state, stopwords=stopwords_set, mask=mask, contour_width=contour_width, contour_color=contour_color, @@ -206,7 +200,10 @@ class ComfyWordCloud: image_colors = ImageColorGenerator(np.array(p_color_ref_image)) wc.recolor(color_func=image_colors) - return (pil2tensor(wc.to_image()),) + ret_image = wc.to_image().convert('RGBA') + ret_mask = pil2tensor(ret_image)[0, :, :, 3] + + return (pil2tensor(ret_image), ret_mask,) NODE_CLASS_MAPPINGS = {