myapp/palette.py

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 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)

    return image


def extract_color_clusters(image_array: np.ndarray | Image.Image, n_clusters=2):
    if not isinstance(image_array, np.ndarray):
        image_array = np.array(image_array)

    w, h, d = image_array.shape
    pixels = image_array.reshape(w * h, d)

    return KMeans(n_clusters=n_clusters).fit(pixels)


def array_to_hex(values: np.ndarray):
    values = np.round(values).astype(int)
    return "#" + ("{:02X}" * len(values)).format(*values)