102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
import torch
|
|
import numpy as np
|
|
from sklearn.cluster import KMeans
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import textwrap
|
|
|
|
class ImageColorPaletteAnalyzer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"num_colors": ("INT", {"default": 5, "min": 2, "max": 10, "step": 1}),
|
|
"image_width": ("INT", {"default": 512, "min": 256, "max": 1024, "step": 64}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("palette_image", "color_info")
|
|
FUNCTION = "analyze_colors"
|
|
CATEGORY = "image/analysis"
|
|
|
|
def analyze_colors(self, image, num_colors, image_width):
|
|
# ComfyUI images are in format [B, H, W, C] and values 0-1
|
|
# We process the first image in the batch
|
|
img_tensor = image[0]
|
|
img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
# Reshape the image to be a list of pixels
|
|
pixels = img_np.reshape(-1, 3)
|
|
|
|
# Use KMeans to find dominant colors
|
|
kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)
|
|
kmeans.fit(pixels)
|
|
|
|
# Get the colors and their counts
|
|
colors = kmeans.cluster_centers_.astype(np.uint8)
|
|
labels = kmeans.labels_
|
|
counts = np.bincount(labels)
|
|
|
|
# Calculate percentages
|
|
total_pixels = len(pixels)
|
|
percentages = (counts / total_pixels) * 100
|
|
|
|
# Sort colors by percentage (descending)
|
|
sorted_indices = np.argsort(counts)[::-1]
|
|
sorted_colors = colors[sorted_indices]
|
|
sorted_percentages = percentages[sorted_indices]
|
|
|
|
# Generate hex codes and format string output
|
|
color_info_lines = []
|
|
hex_colors = []
|
|
for color, pct in zip(sorted_colors, sorted_percentages):
|
|
hex_color = "#{:02x}{:02x}{:02x}".format(color[0], color[1], color[2])
|
|
hex_colors.append(hex_color)
|
|
color_info_lines.append(f"{hex_color}: {pct:.2f}%")
|
|
|
|
color_info_str = "\n".join(color_info_lines)
|
|
|
|
# Create a visual palette image
|
|
height = 150
|
|
palette_img = Image.new('RGB', (image_width, height), color=(255, 255, 255))
|
|
draw = ImageDraw.Draw(palette_img)
|
|
|
|
# Draw color swatches proportional to their percentage
|
|
current_x = 0
|
|
for color, pct, hex_code in zip(sorted_colors, sorted_percentages, hex_colors):
|
|
width = int((pct / 100) * image_width)
|
|
# Ensure the last block fills the remaining space due to rounding
|
|
if hex_code == hex_colors[-1]:
|
|
width = image_width - current_x
|
|
|
|
# Draw color rectangle
|
|
draw.rectangle(
|
|
[current_x, 0, current_x + width, height - 50],
|
|
fill=tuple(color)
|
|
)
|
|
|
|
# Draw hex text (simple drawing, default font)
|
|
# Centered in its block if possible
|
|
text_x = current_x + (width // 2) - 25
|
|
if width > 50: # Only draw text if the block is wide enough
|
|
draw.text((text_x, height - 35), hex_code, fill=(0, 0, 0))
|
|
draw.text((text_x, height - 20), f"{pct:.1f}%", fill=(50, 50, 50))
|
|
|
|
current_x += width
|
|
|
|
# Convert back to ComfyUI tensor format [B, H, W, C]
|
|
palette_np = np.array(palette_img).astype(np.float32) / 255.0
|
|
palette_tensor = torch.from_numpy(palette_np).unsqueeze(0)
|
|
|
|
return (palette_tensor, color_info_str)
|
|
|
|
# Registration dictionaries for ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageColorPaletteAnalyzer": ImageColorPaletteAnalyzer
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageColorPaletteAnalyzer": "🎨 Image Color Palette Analyzer"
|
|
}
|