nodes/chflame163-ComfyUI_WordCloud
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Word Cloud node add mask output

Browse Source
This commit is contained in:
chflame163
2024-01-16 18:18:26 +08:00
committed by GitHub
parent e685edd8a5
commit d8b633ae61
1 changed files with 19 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/comfy_wordcloud.py
+19 -22
View File
@@ -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 = {