Dictionary of Applied Machine Learning · clustering
Numerical companion to the entry clustering: it recomputes what the entry states and prints one line per check
Run it with python3 clustering.py, from any directory — it writes its output files into the current directory. Requires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Download clustering.py · Notebook · Open in Colab
#!/usr/bin/env python3
"""
clustering.py — Foreground/background segmentation of the cow image via
hard clustering (k-means) and soft clustering (Gaussian mixture model).
Pixels of ``assets/CowsAustria.jpg`` are used as a 3-dimensional dataset
in RGB space. Two clusters are fit:
1. k-means (hard assignment: each pixel → one cluster)
2. GMM/EM (soft assignment: each pixel → posterior over clusters)
The script writes two PNGs into ``pythondemos/`` for quick preview and a
single preview PDF showing original / hard / soft side by side.
Run from the repo root:
python3 pythondemos/clustering.py
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.image import imread
REPO = Path(__file__).resolve().parent.parent
IMG_PATH = REPO / "assets" / "BergSee.jpg"
OUT_DIR = REPO / "pythondemos"
# Export the PNG panels at 300 ppi for their printed size. Each panel is
# placed in a 0.31\textwidth minipage (\textwidth = 390pt), so its printed
# width is 0.31 * 390 / 72.27 in; at 300 ppi that is ~502 px. Saving the
# full-resolution source (2816 px wide) embedded ~30x more pixels than the
# page can show and dominated the book PDF size.
TEXTWIDTH_PT = 390.0
PT_PER_INCH = 72.27
PANEL_FRACTION = 0.31
TARGET_DPI = 300
MAX_WIDTH_PX = round(PANEL_FRACTION * TEXTWIDTH_PT / PT_PER_INCH * TARGET_DPI)
RNG = np.random.default_rng(0)
def save_png_300dpi(img01: np.ndarray, path: Path, max_width_px: int = MAX_WIDTH_PX) -> None:
"""Save an (H, W, 3) array in [0,1] as PNG, downscaled so its printed
width (MAX_WIDTH_PX) is TARGET_DPI ppi. Pillow is matplotlib's PNG
backend (imread/imsave already rely on it), so this adds no dependency."""
from PIL import Image
rgb8 = (np.clip(img01, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
im = Image.fromarray(rgb8)
if im.width > max_width_px:
new_h = round(im.height * max_width_px / im.width)
im = im.resize((max_width_px, new_h), Image.LANCZOS)
im.save(path)
def _kmeanspp_init(X: np.ndarray, k: int) -> np.ndarray:
"""k-means++ seeding: spreads initial centroids across the dataset."""
n = len(X)
centroids = [X[RNG.integers(n)]]
for _ in range(k - 1):
d2 = np.min(
((X[:, None, :] - np.stack(centroids)[None, :, :]) ** 2).sum(-1),
axis=1,
)
probs = d2 / d2.sum()
centroids.append(X[RNG.choice(n, p=probs)])
return np.stack(centroids)
def kmeans(X: np.ndarray, k: int, n_iter: int = 20) -> tuple[np.ndarray, np.ndarray]:
"""Lloyd's algorithm with k-means++ seeding."""
centroids = _kmeanspp_init(X, k).copy()
for _ in range(n_iter):
d2 = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
labels = d2.argmin(axis=1)
for c in range(k):
mask = labels == c
if mask.any():
centroids[c] = X[mask].mean(axis=0)
return centroids, labels
def gmm_em(X: np.ndarray, k: int, n_iter: int = 40) -> np.ndarray:
"""Fit a diagonal-covariance GMM via EM. Returns soft assignments (n,k)."""
n, d = X.shape
idx = RNG.choice(n, size=k, replace=False)
mu = X[idx].copy()
var = np.tile(X.var(axis=0) + 1e-3, (k, 1))
pi = np.full(k, 1.0 / k)
for _ in range(n_iter):
# E-step: log N(x | mu_c, diag(var_c))
log_p = np.empty((n, k))
for c in range(k):
diff = X - mu[c]
log_p[:, c] = (
np.log(pi[c] + 1e-12)
- 0.5 * np.sum(np.log(2 * np.pi * var[c]))
- 0.5 * np.sum(diff ** 2 / var[c], axis=1)
)
log_p -= log_p.max(axis=1, keepdims=True)
resp = np.exp(log_p)
resp /= resp.sum(axis=1, keepdims=True)
# M-step
nk = resp.sum(axis=0) + 1e-12
pi = nk / n
mu = (resp.T @ X) / nk[:, None]
for c in range(k):
diff = X - mu[c]
var[c] = (resp[:, c, None] * diff ** 2).sum(axis=0) / nk[c] + 1e-4
return resp
def main() -> None:
img = imread(IMG_PATH).astype(np.float32) / 255.0
H, W, _ = img.shape
# Downsample for speed; label every pixel at the end.
step = 4
small = img[::step, ::step].reshape(-1, 3)
K = 2
# ---- hard clustering ----
centroids, _ = kmeans(small, k=K, n_iter=25)
X = img.reshape(-1, 3)
d2 = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
hard_labels = d2.argmin(axis=1)
# Pick one cluster to highlight: the brightest centroid, which
# corresponds to the sky/mountain region.
target = int(np.argmax(centroids.sum(axis=1)))
# Hard: keep the original RGB for pixels assigned to the target
# cluster; black out the rest.
hard_mask = (hard_labels == target).reshape(H, W, 1).astype(np.float32)
hard_img = img * hard_mask
# ---- soft clustering ----
# Fit a GMM on the full image, then soften the posteriors with a
# temperature so that the visualization shows graded membership
# rather than a near-binary assignment.
# Fit the GMM, then inflate the per-component variances to widen
# each Gaussian and re-score with the posterior formula. Larger
# variances give softer (more uniform) posteriors, producing
# smoother transitions between clusters.
resp = gmm_em(X, k=K, n_iter=25)
# Recover component means and recompute posteriors with inflated
# diagonal covariance (same variance in each component for the
# visualization — we only need softer posteriors, not a better
# fit).
nk = resp.sum(axis=0) + 1e-12
mu = (resp.T @ X) / nk[:, None]
pi = nk / nk.sum()
var = np.full((K, 3), 0.25) # wide, uniform covariance
log_p = np.empty((len(X), K))
for c in range(K):
diff = X - mu[c]
log_p[:, c] = (
np.log(pi[c] + 1e-12)
- 0.5 * np.sum(np.log(2 * np.pi * var[c]))
- 0.5 * np.sum(diff ** 2 / var[c], axis=1)
)
log_p -= log_p.max(axis=1, keepdims=True)
resp = np.exp(log_p)
resp /= resp.sum(axis=1, keepdims=True)
# Align the GMM component order with k-means centroids by matching
# GMM means to the k-means centroids (greedy match).
mu_full = np.stack([
(resp[:, c, None] * X).sum(axis=0)
/ (resp[:, c].sum() + 1e-12)
for c in range(K)
])
perm = []
used = set()
for c in range(K):
order = np.argsort(np.linalg.norm(mu_full - centroids[c], axis=1))
for j in order:
if j not in used:
perm.append(j)
used.add(j)
break
resp = resp[:, perm]
# Soft: modulate each pixel's original RGB by the posterior
# probability that it belongs to the target cluster.
soft_weight = resp[:, target].reshape(H, W, 1)
soft_img = img * soft_weight
save_png_300dpi(img, OUT_DIR / "clustering_original.png")
save_png_300dpi(hard_img, OUT_DIR / "clustering_hard.png")
save_png_300dpi(soft_img, OUT_DIR / "clustering_soft.png")
# ---- preview PDF ----
fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))
axes[0].imshow(img)
axes[0].set_title("original")
axes[1].imshow(np.clip(hard_img, 0, 1))
axes[1].set_title(f"hard: k-means (k={K})")
axes[2].imshow(np.clip(soft_img, 0, 1))
axes[2].set_title(f"soft: GMM posterior (k={K})")
for ax in axes:
ax.set_xticks([])
ax.set_yticks([])
fig.tight_layout()
fig.savefig(OUT_DIR / "clustering.png", dpi=110)
print(f"wrote {OUT_DIR/'clustering.png'}")
print(f"wrote {OUT_DIR/'clustering_original.png'}")
print(f"wrote {OUT_DIR/'clustering_hard.png'}")
print(f"wrote {OUT_DIR/'clustering_soft.png'}")
if __name__ == "__main__":
main()