color quantization using gaussian mixture model in OKLCH

gmm.py

import numpy as np
from PIL import Image
from sklearn.mixture import GaussianMixture
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import ColorMath

IMAGE_PATH  = "data/adirondack_chairs.png"
OUTPUT_PATH = "output/quant_oklch.png"
K           = 10
N_GROUPS    = 10
SIGMA_SCALE = 0.5  # 1.0 = full spread, 0.0 = always the mean

HUE_PERIOD = 2 * np.pi

# ── load & convert to OKLCH ────────────────────────────────────────
img    = Image.open(IMAGE_PATH).convert("RGB")
pixels = np.array(img).reshape(-1, 3)

rgb_ints = ((pixels[:, 0].astype(np.uint32) << 16)
          | (pixels[:, 1].astype(np.uint32) << 8)
          |  pixels[:, 2].astype(np.uint32))

oklab = ColorMath.rgb_ints_to_oklab(rgb_ints)

# Convert each OKLab pixel to OKLCH
oklch = ColorMath.oklab_to_oklch(oklab).astype(np.float32)

# ── fit GMM in OKLCH ───────────────────────────────────────────────
gmm = GaussianMixture(n_components=K, covariance_type="full", random_state=42)
gmm.fit(oklch)

# ── sample & plot ──────────────────────────────────────────────────
fig, axes = plt.subplots(
    N_GROUPS, K,
    figsize=(K * 1.2, N_GROUPS * 0.7),
    gridspec_kw={"hspace": 0.35, "wspace": 0.0}
)

rng = np.random.default_rng()

for row in range(N_GROUPS):
    for col in range(K):
        # sample one colour directly from component col's Gaussian in OKLCH
        lch = rng.multivariate_normal(
            gmm.means_[col],
            gmm.covariances_[col] * SIGMA_SCALE**2
        )

        # clamp / wrap into a sane OKLCH range
        lch[0] = np.clip(lch[0], 0.0, 1.0)   # L
        lch[1] = max(lch[1], 0.0)            # C
        lch[2] = np.mod(lch[2], HUE_PERIOD)  # h

        # convert back to sRGB for display
        lab  = ColorMath.oklch_to_oklab(lch)
        srgb = np.clip(ColorMath.oklab_to_srgb(lab), 0, 1)

        ax = axes[row, col]
        ax.add_patch(patches.Rectangle((0, 0), 1, 1, color=srgb))
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis("off")

fig.suptitle(f"GMM palette samples in OKLCH  (K={K}, {N_GROUPS} draws)", fontsize=11, y=1.01)
plt.savefig(OUTPUT_PATH, dpi=150, bbox_inches="tight")
print(f"saved → {OUTPUT_PATH}")