diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8d35cb3 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__ +*.pyc diff --git a/README.md b/README.md new file mode 100644 index 0000000..31a831d --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +#ComfyUI-ColorWheel + +I made this to show the dominant colors in an image on a color wheel. + +The image to color node is from [https://github.com/JayLyu/ComfyUI_BaiKong_Node](https://github.com/JayLyu/ComfyUI_BaiKong_Node) +I copied it over here because it's the only node from the pack needed for what this node was designed for. + +What is was designed for is color studies, particularly for sunsets but it will work for anything. The workflow it's used for is in +the file colorstudy-workflow.json. It looks like this: diff --git a/colorstudy-workflow.json b/colorstudy-workflow.json new file mode 100644 index 0000000..f4452a1 --- /dev/null +++ b/colorstudy-workflow.json @@ -0,0 +1,635 @@ +{ + "last_node_id": 21, + "last_link_id": 31, + "nodes": [ + { + "id": 8, + "type": "PreviewImage", + "pos": { + "0": 1951.1346435546875, + "1": -221.58975219726562 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 18 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 6, + "type": "PreviewImage", + "pos": { + "0": 1711.134765625, + "1": -221.58975219726562 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 3, + "type": "PreviewImage", + "pos": { + "0": 1711.134765625, + "1": 78.41023254394531 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 2 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 19, + "type": "AccurateColorWheelNode", + "pos": { + "0": 939.3956298828125, + "1": 229.64263916015625 + }, + "size": { + "0": 290.7735290527344, + "1": 69.09516143798828 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "color_string", + "type": "STRING", + "link": 27, + "widget": { + "name": "color_string" + } + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 31 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "AccurateColorWheelNode" + }, + "widgets_values": [ + "#beb6bb,#c38064,#9a99ac,#e9b68b,#ead5c2", + 512 + ] + }, + { + "id": 12, + "type": "BK_Img2Color", + "pos": { + "0": 589.3956298828125, + "1": 82.73787689208984 + }, + "size": { + "0": 315, + "1": 216 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "input_image", + "type": "IMAGE", + "link": 21 + } + ], + "outputs": [ + { + "name": "COLORS", + "type": "STRING", + "links": [ + 24 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "SELECT_COLOR", + "type": "STRING", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BK_Img2Color" + }, + "widgets_values": [ + 6, + false, + 80, + "#feffff", + 1, + "#fefeff, #f7ad89, #c5c1ce, #fbca57, #d1adae, #f2c2ae,#fefeff" + ] + }, + { + "id": 7, + "type": "MaskToImage", + "pos": { + "0": 939.3956298828125, + "1": -187.26214599609375 + }, + "size": { + "0": 293.2176513671875, + "1": 26.86839485168457 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "mask", + "type": "MASK", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "MaskToImage" + }, + "widgets_values": [] + }, + { + "id": 2, + "type": "ColorPreprocessor", + "pos": { + "0": 939.3956298828125, + "1": -27.262130737304688 + }, + "size": { + "0": 292.6774597167969, + "1": 58 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 4 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ColorPreprocessor" + }, + "widgets_values": [ + 1024 + ] + }, + { + "id": 10, + "type": "LayerUtility: ImageAutoCrop V2", + "pos": { + "0": 179.3956298828125, + "1": -187.26214599609375 + }, + "size": { + "0": 378, + "1": 434 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 16 + }, + { + "name": "mask", + "type": "MASK", + "link": 17, + "shape": 7 + } + ], + "outputs": [ + { + "name": "cropped_image", + "type": "IMAGE", + "links": [ + 18, + 19, + 21 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "box_preview", + "type": "IMAGE", + "links": null, + "shape": 3 + }, + { + "name": "cropped_mask", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LayerUtility: ImageAutoCrop V2" + }, + "widgets_values": [ + true, + "#FFFFFF", + "1:1", + 1, + 1, + "None", + 1024, + "min_bounding_rect", + 100, + 0, + "RMBG 1.4", + "sam_vit_h (2.56GB)", + "GroundingDINO_SwinT_OGC (694MB)", + 0.3, + "subject" + ], + "color": "rgba(38, 73, 116, 0.7)" + }, + { + "id": 4, + "type": "ImageResize+", + "pos": { + "0": 589.3956298828125, + "1": -187.26214599609375 + }, + "size": { + "0": 315, + "1": 218 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 19 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "width", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "height", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 1024, + 1024, + "nearest", + "keep proportion", + "always", + 0 + ] + }, + { + "id": 16, + "type": "ShowText|pysssss", + "pos": { + "0": 939.3956298828125, + "1": 82.73787689208984 + }, + "size": { + "0": 290.13885498046875, + "1": 76 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "text", + "type": "STRING", + "link": 24, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 27 + ], + "slot_index": 0, + "shape": 6 + } + ], + "properties": { + "Node name for S&R": "ShowText|pysssss" + }, + "widgets_values": [ + "", + "#fefeff, #f7ad89, #c5c1ce, #fbca57, #d1adae, #f2c2ae" + ] + }, + { + "id": 18, + "type": "PreviewImage", + "pos": { + "0": 1951, + "1": 78 + }, + "size": { + "0": 210, + "1": 246 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 31 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 1, + "type": "LoadImage", + "pos": { + "0": 1281.134765625, + "1": -221.58975219726562 + }, + "size": { + "0": 401.2361755371094, + "1": 540.2031860351562 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 16 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "MASK", + "type": "MASK", + "links": [ + 8, + 17 + ], + "slot_index": 1, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "clipspace-mask-200317.png [temp]", + "image" + ] + } + ], + "links": [ + [ + 2, + 2, + 0, + 3, + 0, + "IMAGE" + ], + [ + 4, + 4, + 0, + 2, + 0, + "IMAGE" + ], + [ + 8, + 1, + 1, + 7, + 0, + "MASK" + ], + [ + 9, + 7, + 0, + 6, + 0, + "IMAGE" + ], + [ + 16, + 1, + 0, + 10, + 0, + "IMAGE" + ], + [ + 17, + 1, + 1, + 10, + 1, + "MASK" + ], + [ + 18, + 10, + 0, + 8, + 0, + "IMAGE" + ], + [ + 19, + 10, + 0, + 4, + 0, + "IMAGE" + ], + [ + 21, + 10, + 0, + 12, + 0, + "IMAGE" + ], + [ + 24, + 12, + 0, + 16, + 0, + "STRING" + ], + [ + 27, + 16, + 0, + 19, + 0, + "STRING" + ], + [ + 31, + 19, + 0, + 18, + 0, + "IMAGE" + ] + ], + "groups": [ + { + "title": "Group", + "bounding": [ + 169, + -261, + 1073, + 570 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "title": "Group", + "bounding": [ + 1271, + -296, + 900, + 630 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 1.4122927695244532, + "offset": [ + -1208.2183921954272, + 321.92843210193087 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py index ca0bc07..96056f0 100644 --- a/nodes.py +++ b/nodes.py @@ -4,16 +4,29 @@ from PIL import Image, ImageDraw, ImageFont import colorsys import math -class AccurateColorWheelNode: +from typing import Tuple, List +import sys +import os +import torch + +from sklearn.cluster import KMeans + +from numpy import ndarray +import numpy as np + +sys.path.append(os.path.dirname(os.path.dirname(__file__))) + + +class ColorWheelNode: @classmethod def INPUT_TYPES(s): return { "required": { - "color_string": ("STRING", {"default": "#beb6bb,#c38064,#9a99ac,#e9b68b,#ead5c2"}), + "color_string": ("STRING", {"default": "#ead5c2,#9a99ac,#e9b68b"}), "size": ("INT", {"default": 512, "min": 128, "max": 2048}), }, } - + RETURN_TYPES = ("IMAGE",) FUNCTION = "create_accurate_color_wheel" CATEGORY = "image/generation" @@ -39,55 +52,165 @@ class AccurateColorWheelNode: saturation = distance / radius r, g, b = colorsys.hsv_to_rgb(angle, saturation, 1) h, s, v = colorsys.rgb_to_hsv(r, g, b) - - # Convert HSV to 2D vector + vector_x = s * math.cos(h * 2 * math.pi) vector_y = s * math.sin(h * 2 * math.pi) - - # Map to image coordinates + img_x = int((vector_x + 1) * center) img_y = int((vector_y + 1) * center) - + if 0 <= img_x < size and 0 <= img_y < size: if v > value_map[img_y, img_x]: output[img_y, img_x] = [r, g, b] value_map[img_y, img_x] = v - # Create PIL Image for drawing img = Image.fromarray((output * 255).astype(np.uint8)) draw = ImageDraw.Draw(img) - # Place color markers and labels for color in colors: rgb = self.hex_to_rgb(color) h, s, v = colorsys.rgb_to_hsv(rgb[0]/255, rgb[1]/255, rgb[2]/255) - + angle = h * 2 * math.pi distance = s * radius x = center + int(distance * math.cos(angle)) y = center + int(distance * math.sin(angle)) - # Draw color circle - circle_radius = int(size * 0.03) + circle_radius = int(size * 0.02) draw.ellipse((x - circle_radius, y - circle_radius, x + circle_radius, y + circle_radius), fill=color, outline='white') - # Add color label - label_distance = radius + int(size * 0.08) - label_x = center + int(label_distance * math.cos(angle)) - label_y = center + int(label_distance * math.sin(angle)) - draw.text((label_x, label_y), color, fill='black', anchor='mm', stroke_width=2, stroke_fill='white') + # Add palette block below the wheel + palette_height = int(size * 0.2) # Reduced height + palette_width = size + palette_block = Image.new('RGB', (palette_width, palette_height), color='white') + palette_draw = ImageDraw.Draw(palette_block) - # Convert to numpy array and then to PyTorch tensor - np_image = np.array(img).astype(np.float32) / 255.0 + font_size = int(size * 0.025) + try: + font = ImageFont.truetype("arial.ttf", font_size) + except IOError: + font = ImageFont.load_default() + + block_size = int(palette_height * 0.5) # Smaller block size + block_spacing = int(palette_height * 0.25) + for i, color in enumerate(colors): + block_x = i * (block_size + block_spacing) + block_spacing + block_y = int(palette_height * 0.25) + palette_draw.rectangle((block_x, block_y, block_x + block_size, block_y + block_size), fill=color) + + label_x = block_x + label_y = block_y + block_size + int(palette_height * 0.05) + palette_draw.text((label_x, label_y), color, fill='black', font=font) + + combined_img = Image.new('RGB', (size, size + palette_height), color='white') + combined_img.paste(img, (0, 0)) + combined_img.paste(palette_block, (0, size)) + + np_image = np.array(combined_img).astype(np.float32) / 255.0 torch_image = torch.from_numpy(np_image)[None,] return (torch_image,) +class BK_Img2Color: + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_image": ("IMAGE",), + }, + "optional": { + "num_colors": ("INT", {"default": 1, "min": 1, }), + "get_complementary_color": ( + "BOOLEAN", + { + "default": False, + "label_off": "false", + "label_on": "true", + }, + ), + "accuracy": ( + "INT", + { + "default": 80, + "display": "slider", + "min": 1, + "max": 100, + }, + ), + "exclude_colors": ( + "STRING", + { + "default": "", + }, + ), + "select_color": ("INT", { + "default": 1, + "min": 1, + }), + } + } + + RETURN_TYPES = ("STRING", "STRING",) + RETURN_NAMES = ("COLORS", "SELECT_COLOR",) + CATEGORY = "⭐️ Baikong/Color" + FUNCTION = "main" + # OUTPUT_NODE = True + DESCRIPTION = "从输入图像中提取主要颜色，可指定颜色数量，支持排除特定颜色，并可选择生成互补色" + + def __init__(self): + pass + + def main(self, input_image: torch.Tensor, num_colors: int = 5, accuracy: int = 80, + get_complementary_color: bool = False, exclude_colors: str = "", select_color: int = 1) -> Tuple[str, str]: + self.exclude = [color.strip().lower() for color in exclude_colors.strip().split( + ",")] if exclude_colors.strip() else [] + self.num_iterations = int(512 * (accuracy / 100)) + + original_colors = self.interrogate_colors(input_image, num_colors) + rgb = self.ndarrays_to_rgb(original_colors) + + if get_complementary_color: + rgb = self.rgb_to_complementary(rgb) + + hex_colors = [ + f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}" for color in rgb] + out = self.join_and_exclude(hex_colors) + + # 处理 select_color + color_list = out.split(", ") + selected_color = color_list[-1] if select_color > len( + color_list) else color_list[select_color - 1] + + # 指定的输出格式 {"ui": {"text": (value1, value2)}, "result": (value1, value2)} + return {"ui": {"text": (out, selected_color)}, "result": (out, selected_color)} + + def join_and_exclude(self, colors: List[str]) -> str: + return ", ".join( + [str(color) + for color in colors if color.lower() not in self.exclude] + ) + + def rgb_to_complementary( + self, colors: List[Tuple[int, int, int]] + ) -> List[Tuple[int, int, int]]: + return [(255 - color[0], 255 - color[1], 255 - color[2]) for color in colors] + + def ndarrays_to_rgb(self, colors: List[ndarray]) -> List[Tuple[int, int, int]]: + return [(int(color[0]), int(color[1]), int(color[2])) for color in colors] + + def interrogate_colors(self, image: torch.Tensor, num_colors: int) -> List[ndarray]: + pixels = image.view(-1, image.shape[-1]).numpy() + kmeans = KMeans(n_clusters=num_colors, algorithm="lloyd", + max_iter=self.num_iterations, n_init=10) + colors = kmeans.fit(pixels).cluster_centers_ * 255 + return colors + NODE_CLASS_MAPPINGS = { - "AccurateColorWheelNode": AccurateColorWheelNode + "AccurateColorWheelNode": ColorWheelNode } NODE_DISPLAY_NAME_MAPPINGS = { - "AccurateColorWheelNode": "Accurate Color Wheel Generator" + "AccurateColorWheelNode": "Color Wheel Generator" } diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ba92068 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +torch==2.4.1 +numpy==1.26.4 +pillow==10.4.0 +scikit-learn==1.5.2