K-Means Color Quantization

K-means color quantization treats each pixel color as a point in 3D space (R, G, B) and groups these points into K clusters so that colors inside each cluster are similar. The cluster centers become a palette of K representative colors, and each pixel is replaced by the color of its cluster’s center. This reduces the number of distinct colors while trying to keep the image looking similar to the original.

Mathematically, K-means tries to minimize the sum of squared distances from each pixel to its assigned center in RGB space. Using Euclidean distance, every pixel is assigned to the nearest center, then centers are recomputed as the mean of all pixels that chose that center. Repeating this alternation of “assign” and “update” steps causes centers to move toward dense regions in color space and eventually stabilize (converge).

---

1. Background Knowledge

• Color as a 3D vector: An RGB pixel can be written as a vector $p = [r, g, b]$ in $\mathbb{R}^3$. Color quantization is then just vector quantization in 3D: approximating many vectors with a smaller set of prototype vectors (the centers).

• K-means clustering basics: K-means is an unsupervised clustering algorithm. Given data points $x_1, \dots, x_N$ and a target number of clusters $K$, it alternates between:

• Assignment step: assign each point to its nearest center.

• Update step: recompute each center as the mean of its assigned points.

It stops when assignments stop changing or when centers move less than a tolerance.

• Objective function: K-means minimizes:

$$J = \sum_{i=1}^{N} \min_{1 \le k \le K} \|x_i - \mu_k\|^2$$

where $\mu_k$ are the centers. For color quantization, each $x_i$ is a pixel, and $\mu_k$ is a palette color.

---

2. Algorithm / Approach Pattern

General approach for K-means color quantization:

1. Flatten the image into a 2D array of pixels: shape $(N, 3)$ where $N = H \times W$.
2. Initialize K centers (e.g., choose K random pixels).
3. Repeat until convergence or max iterations:

• Compute distances from all pixels to all centers.
• Assign each pixel to the nearest center.
• Update each center as the mean color of the pixels assigned to it.

4. Reconstruct the output image by replacing every pixel’s color with its center’s color and reshaping back to $(H, W, 3)$.

This is the standard batch K-means pattern, specialized to 3D RGB data.

---

3. Step-by-Step Strategy (Implementation Outline)

Assume you have:

• Input: image as a NumPy array of shape (H, W, 3) and integer K.

Step 1: Preprocessing

• Convert to float for arithmetic (if not already):

img = image.astype(np.float32)

• Reshape to pixel list:

H, W, C = img.shape # C should be 3
pixels = img.reshape(-1, 3) # shape (N, 3), N = H*W
N = pixels.shape

Step 2: Initialize Cluster Centers

• Simple approach: random subset of pixels:

rng = np.random.default_rng(seed)
indices = rng.choice(N, size=K, replace=False)
centers = pixels[indices] # shape (K, 3)

Step 3: Main K-means Loop

Repeat for a fixed number of iterations or until convergence:

1. Compute distances and assignments

Efficient vectorized distance computation:

# pixels: (N, 3), centers: (K, 3)
# expand and use broadcasting
diff = pixels[:, None, :] - centers[None, :, :] # (N, K, 3)
dists_sq = np.sum(diff**2, axis=2) # (N, K)
labels = np.argmin(dists_sq, axis=1) # (N,), cluster index for each pixel

2. Update centers

For each k in 0..K-1, compute mean of assigned pixels:

new_centers = np.zeros_like(centers)
for k in range(K):
mask = (labels == k)
if np.any(mask):
new_centers[k] = pixels[mask].mean(axis=0)
else:
# handle empty cluster (e.g., reinitialize randomly)
new_centers[k] = pixels[rng.integers(N)]

3. Check convergence

shift = np.linalg.norm(new_centers - centers)
if shift < tol:
break
centers = new_centers

Step 4: Reconstruct Quantized Image

• Map each pixel to its center color and reshape:

quantized_pixels = centers[labels] # shape (N, 3)
quantized_img = quantized_pixels.reshape(H, W, 3)
# Optionally cast back to uint8
quantized_img = np.clip(quantized_img, 0, 255).astype(np.uint8)

---

4. Common Pitfalls

• Incorrect reshaping:

• Forgetting to keep the last dimension as 3 (RGB).

• Mixing up (H, W, 3) and (N, 3) shapes.

• Inefficient loops:

• Computing distance from each pixel to each center using pure Python loops (for over pixels or centers) will be very slow.

• Use NumPy broadcasting for distances and assignments.

• Empty clusters:

• Some clusters may get no assigned points, leading to mean of an empty slice (NaNs).

• You must detect this and reinitialize such centers (e.g., with a random pixel).

• Data types and overflow:

• Distance computation on uint8 can overflow: do arithmetic in float32 or float64.

• Remember to cast back to uint8 for image display/saving.

• Non-convergence conditions:

• If you only rely on a tolerance and it’s too strict, you may run many unnecessary iterations.

• Use a max iterations cap (e.g., 10–20) plus a tolerance.

---

5. Time & Space Complexity

Let:

• $N = H \times W$ (number of pixels),

• $K$ = number of clusters,

• $T$ = number of K-means iterations.

• Time Complexity:

• Each iteration:

• Distance computation: $O(N \times K)$

• Assignment: $O(N \times K)$ (argmin over K for each pixel)

• Center update: $O(N)$ to aggregate pixels into K clusters.

• Total: $O(T \times N \times K)$

• Space Complexity:

• Storing pixels: $O(N)$ for (N, 3)

• Centers: $O(K)$

• Distance matrix (if stored fully): $O(N \times K)$ (though you can sometimes reduce memory by computing in chunks).

• Overall: typically $O(N + K + N K)$; in practice dominated by N*K if K is not tiny and you store all distances at once.