add generate palette command

myapp/cli.py

@@ -2,6 +2,7 @@ import click
 import dotenv
 from huggingface_hub import InferenceClient
 
+from myapp.palette import extract_color_clusters, generate_palette_image
 from myapp.types import ImageParamType
 
 dotenv.load_dotenv()
@@ -32,5 +33,10 @@ def generate_image(prompt, target, model, width, height):
 
 @cli.command()
 @click.option("--image", type=ImageParamType(), required=True)
-def generate_palette(image):
-    pass
+@click.option("--target", type=click.Path(dir_okay=False), required=True)
+@click.option("--n-colors", default=4)
+@click.option("--shade", "shades", default=(0.0,), multiple=True)
+def generate_palette(image, target, n_colors, shades):
+    k_means = extract_color_clusters(image, n_colors)
+    palette = generate_palette_image(k_means, shades=shades)
+    palette.save(target)

myapp/palette.py

@@ -1,25 +1,20 @@
+import colorsys
+import itertools
+
 import numpy as np
 from PIL import Image
 from sklearn.cluster import KMeans
 
 
-def join_images(a: Image.Image, b: Image.Image):
-    result = Image.new(a.mode, (a.width + b.width, max(a.height, b.height)))
-    result.paste(a)
-    result.paste(b, (a.width, 0))
-
-    return result
-
+def get_hsv_value(cluster: np.ndarray):
+    return colorsys.rgb_to_hsv(*cluster / 255)[2]
 
-def generate_palette_image(model: KMeans, size=40):
-    image = Image.new("RGB", (0, size))
 
-    for cluster_center in model.cluster_centers_:
-        color = tuple(map(int, cluster_center))
-        part = Image.new("RGB", (40, 40), color)
-        image = join_images(image, part)
+def add_hsv_saturation(cluster: np.ndarray, delta: float):
+    h, s, v = colorsys.rgb_to_hsv(*cluster / 255)
+    s = max(0, min(1, s + delta))
 
-    return image
+    return np.array(colorsys.hsv_to_rgb(h, s, v)) * 255
 
 
 def extract_color_clusters(image_array: np.ndarray | Image.Image, n_clusters=2):
@@ -32,6 +27,26 @@ def extract_color_clusters(image_array: np.ndarray | Image.Image, n_clusters=2):
     return KMeans(n_clusters=n_clusters).fit(pixels)
 
 
+def iter_color_shades(k_means: KMeans, shades: tuple[float, ...]):
+    cluster_centers = sorted(k_means.cluster_centers_, key=get_hsv_value)
+
+    for delta, cluster_center in itertools.product(shades, cluster_centers):
+        yield add_hsv_saturation(cluster_center, delta)
+
+
+def generate_palette_image(k_means: KMeans, size=40, shades=(0.0,)):
+    num_cluster_centers = len(k_means.cluster_centers_)
+    image = Image.new("RGB", (num_cluster_centers * size, len(shades) * size))
+
+    for i, color in enumerate(iter_color_shades(k_means, shades)):
+        color = tuple(map(int, color))
+        part = Image.new("RGB", (size, size), color)
+        position = (i % num_cluster_centers * size, i // num_cluster_centers * size)
+        image.paste(part, position)
+
+    return image
+
+
 def array_to_hex(values: np.ndarray):
     values = np.round(values).astype(int)
     return "#" + ("{:02X}" * len(values)).format(*values)