K-Means for image compression

I recently finished Andrew Ng’s CS229 course remotely. The course was extremely challenging and covered a wide range of Machine Learning concepts. Even if you have worked with, and used Machine Learning algorithms, he introduced concepts in novel and interesting ways I didn’t expect. One of the things that struct a chord with me was how we worked on K-Means clustering.

The algorithm is quite simple. Courtesy of CS229…

Given a training set, if we want to group it into n clusters

1. Initialize cluster centroids randomly

2. Repeat until convergence

For every i set

For each j set

Here an example of running K-means on the Iris dataset.

But an alternate use for K-Means is image compression. Given an image with the standard RGB colors we have 256x256x256=16.77M colors. We can use K-Means to compress these into less colors. The idea is similar to the above, with a few differences.

1. We take an image of dimension (m,n,3) (3=R,G,B) and resize it to (mxn,3)
2. For k clusters, we randomly sample k points from this data
3. We assign colors closest to the k points we’ve selected to that cluster
4. The new cluster points are the mean R,G,B points for each cluster
5. We repeat this process until convergence (points don’t change, or we reach a threshold)

This is what the code looks like

We can run this on a few images, reducing them from 16.77m colors to 16 colors to see how good the compression is