161 lines
5.8 KiB
Python
161 lines
5.8 KiB
Python
"""
|
|
Dynamic Color Family Grid Extractor (Unsupervised Clustering).
|
|
|
|
Dynamically discovers color families in your image using K-Means in LAB space,
|
|
retains the ORIGINAL image colors for each group, and places them into a single
|
|
grid with separation lines.
|
|
|
|
Requirements:
|
|
pip install numpy pillow scikit-learn --break-system-packages
|
|
|
|
Usage:
|
|
python dynamic_color_families.py input.png --k 5 --out grid_dynamic.png
|
|
"""
|
|
|
|
import argparse
|
|
import math
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from sklearn.cluster import KMeans
|
|
|
|
|
|
def rgb_to_lab(rgb_array):
|
|
"""
|
|
Converts RGB array [0-255] to CIELAB space for perceptual color clustering.
|
|
CIELAB separates lightness (L) from color channels (A, B).
|
|
"""
|
|
# Standard sRGB to XYZ conversion
|
|
rgb = rgb_array.astype(np.float64) / 255.0
|
|
mask = rgb > 0.04045
|
|
rgb[mask] = np.power((rgb[mask] + 0.055) / 1.055, 2.4)
|
|
rgb[~mask] = rgb[~mask] / 12.92
|
|
|
|
# sRGB matrix transformation to D65 XYZ
|
|
transform = np.array([
|
|
[0.4124564, 0.3575761, 0.1804375],
|
|
[0.2126729, 0.7151522, 0.0721750],
|
|
[0.0193339, 0.1191920, 0.9503041]
|
|
])
|
|
xyz = np.dot(rgb, transform.T)
|
|
|
|
# Reference white point D65
|
|
xyz[:, 0] /= 0.95047
|
|
xyz[:, 1] /= 1.00000
|
|
xyz[:, 2] /= 1.08883
|
|
|
|
mask = xyz > 0.008856
|
|
xyz[mask] = np.power(xyz[mask], 1.0 / 3.0)
|
|
xyz[~mask] = (7.787 * xyz[~mask]) + (16.0 / 116.0)
|
|
|
|
L = (116.0 * xyz[:, 1]) - 16.0
|
|
A = 500.0 * (xyz[:, 0] - xyz[:, 1])
|
|
B = 200.0 * (xyz[:, 1] - xyz[:, 2])
|
|
|
|
return np.stack([L, A, B], axis=1)
|
|
|
|
|
|
def create_dynamic_family_grid(image_path, output_path, n_clusters=5, line_width=6, line_color=(180, 180, 180, 255)):
|
|
# Open image with transparency
|
|
img = Image.open(image_path).convert("RGBA")
|
|
arr = np.array(img)
|
|
|
|
h, w, _ = arr.shape
|
|
alpha = arr[:, :, 3]
|
|
fg_mask = alpha > 0 # Ignore background
|
|
|
|
rgb_flat = arr[:, :, :3][fg_mask]
|
|
if len(rgb_flat) == 0:
|
|
raise ValueError("No non-transparent foreground pixels found in image!")
|
|
|
|
print(f"Extracting {n_clusters} dynamic color families using LAB clustering...")
|
|
|
|
# 1. Convert to CIELAB color space (perceptually uniform)
|
|
lab_pixels = rgb_to_lab(rgb_flat)
|
|
|
|
# 2. Dynamically cluster colors into K groups
|
|
# We weight color channels (A, B) slightly higher than Lightness (L) so light/dark shadows
|
|
# of the same color family stay clustered together better.
|
|
features = lab_pixels.copy()
|
|
features[:, 0] *= 0.6 # Scale down Lightness influence to resist shadows
|
|
|
|
kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
|
|
labels = kmeans.fit_predict(features)
|
|
|
|
# Calculate average RGB color for each dynamically found group (for label tags)
|
|
cluster_avg_colors = []
|
|
for i in range(n_clusters):
|
|
avg_rgb = rgb_flat[labels == i].mean(axis=0).astype(int)
|
|
cluster_avg_colors.append(avg_rgb)
|
|
|
|
# Order families by pixel count (largest family first)
|
|
unique_labels, counts = np.unique(labels, return_counts=True)
|
|
sorted_indices = np.argsort(-counts)
|
|
|
|
# 3. Create Grid canvas with separation lines
|
|
cols = math.ceil(math.sqrt(n_clusters))
|
|
rows = math.ceil(n_clusters / cols)
|
|
|
|
grid_w = cols * w + (cols + 1) * line_width
|
|
grid_h = rows * h + (rows + 1) * line_width
|
|
|
|
canvas = Image.new("RGBA", (grid_w, grid_h), (25, 25, 25, 255))
|
|
draw = ImageDraw.Draw(canvas)
|
|
|
|
for rank, cluster_idx in enumerate(sorted_indices):
|
|
r_idx = rank // cols
|
|
c_idx = rank % cols
|
|
|
|
x_start = line_width + c_idx * (w + line_width)
|
|
y_start = line_width + r_idx * (h + line_width)
|
|
|
|
# Mask pixels belonging to this dynamic family
|
|
family_mask_1d = labels == cluster_idx
|
|
full_mask = np.zeros((h, w), dtype=bool)
|
|
full_mask[fg_mask] = family_mask_1d
|
|
|
|
# Extract tile retaining ORIGINAL image pixels
|
|
tile_arr = np.zeros_like(arr)
|
|
tile_arr[full_mask] = arr[full_mask]
|
|
|
|
tile_img = Image.fromarray(tile_arr, mode="RGBA")
|
|
canvas.paste(tile_img, (x_start, y_start), tile_img)
|
|
|
|
# Draw separation lines around panel
|
|
box_coords = [
|
|
x_start - line_width // 2,
|
|
y_start - line_width // 2,
|
|
x_start + w + line_width // 2,
|
|
y_start + h + line_width // 2,
|
|
]
|
|
draw.rectangle(box_coords, outline=line_color, width=line_width)
|
|
|
|
# Label tile with family ID and average RGB swatch
|
|
avg_c = cluster_avg_colors[cluster_idx]
|
|
label_text = f"Family #{rank + 1} ({counts[cluster_idx]} px)"
|
|
|
|
# Label box background
|
|
draw.rectangle(
|
|
[x_start + 10, y_start + 10, x_start + 200, y_start + 40],
|
|
fill=(0, 0, 0, 180)
|
|
)
|
|
# Average color swatch tag
|
|
draw.rectangle(
|
|
[x_start + 15, y_start + 18, x_start + 30, y_start + 33],
|
|
fill=(avg_c[0], avg_c[1], avg_c[2], 255),
|
|
outline=(255, 255, 255, 255)
|
|
)
|
|
draw.text((x_start + 38, y_start + 18), label_text, fill=(255, 255, 255, 255))
|
|
|
|
canvas.save(output_path)
|
|
print(f"Saved dynamic grid with {n_clusters} families to: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Dynamic color family clustering in single grid with separation lines.")
|
|
parser.add_argument("image_path", type=str, help="Path to input image.")
|
|
parser.add_argument("--out", type=str, default="dynamic_families_grid.png", help="Output path for grid image.")
|
|
parser.add_argument("--k", type=int, default=5, help="Number of dynamic color families to discover (default: 5).")
|
|
parser.add_argument("--line-width", type=int, default=6, help="Width of separation lines.")
|
|
args = parser.parse_args()
|
|
|
|
create_dynamic_family_grid(args.image_path, args.out, n_clusters=args.k, line_width=args.line_width) |