Using K-means for image compression

September 19, 2014 - machine learning, tutorial, Octave

I recently used the K-means algorithm to cluster the colors I had in a picture to create a, what I think is, pretty cool profile picture. In this post, I will show you how I managed to do that.

What you will need is a not too large picture (otherwise the algorithm will take much too long) and Octave installed which is available on pretty much any Linux distributions.

The code is available on GitHub with comments and line by line explanations.

Running K-means

I will split the process of running K-means in two different functions:

If you’re familiar with the K-means algorithm you know that the algorithm is based on two steps and so runKMeans will be calling two different functions:

The initialization function

As previously mentioned, initCentroids will pick K (here representing the resulting number of colors we want) random examples from our dataset and choose them as our initial centroids.

It is defined as follows:

function centroids = initCentroids(X, K)

centroids = zeros(K, size(X, 2));
randIndex = randperm(size(X, 1));
centroids = X(randIndex(1:K), :);


The centroid assignment function

assignCentroids loops over every example in the dataset and associates the closest centroid to each and everyone of them.

function idx = assignCentroids(X, centroids)

K = size(centroids, 1);
idx = zeros(size(X, 1), 1);

for i = 1:size(X, 1)
distances = zeros(1, K);

for j = 1:K
distances[1, j] = sumsq(X(i, :) - centroids(j, :));

[one, idx(i)] = min(distances);


The compute centroids function

computeCentroids recomputes the centroids as the mean of the points which were assigned to each centroid previously thanks to assignCentroids.

function centroids = computeCentroids(X, idx, K)

n = size(X, 2);
centroids = zeros(K, n);

for i = 1:K
centroids(i, :) = mean(X(find(idx == i), :));


The function running K-means

runKMeans will call assignCentroids and computeCentroids successively a fixed number of times (thanks to the iter parameter) and will use the initial centroids computed by initCentroids as a starting point.

function [centroids, idx] = runKMeans(X, initialCentroids, iter)

m = size(X, 1);
K = size(initialCentroids, 1);
centroids = initialCentroids;
idx = zeros(m, 1);

for i = 1:iter
idx = assignCentroids(X, centroids);
centroids = computeCentroids(X, idx, K);


Putting things together

In order to tie everything together I wrote a simple script which runs the algorithm and saves the compressed image. It is defined as follows:

#!/usr/bin/octave -q

originalImgPath = argv(){1};
newImgPath = argv(){2};

K = str2num(argv(){3};
iter = str2num(argv(){4};

img = double(imread(originalImgPath));
img = img / 255;
imgSize = size(img);

X = reshape(img, imgSize(1) * imgSize(2), 3);

initialCentroids = initCentroids(X, K);

[centroids, idx] = runKMeans(X, initialCentroids, iter);

idx = assignCentroids(X, centroids);
compressedImg = centroids(idx, :);
compressedImg = reshape(compressedImg, imgSize(1), imgSize(2), 3);

imwrite(compressedImg, newImgPath);

You can also find this script on GitHub.

As you may have noticed, this script takes the following command line arguments:

As an example, this is what I originally had as a picture:

I used the script like so:

./main.m ~/Documents/originalPicture.jpg ~/Documents/newPicture.jpg 8 10

To get back the following image:

Be careful though because, as you might have caught on, the complexity of the whole algorithm is O(i * np * K) where:

So if you have a very large image, or, alternatively if you want a large number of colors or want to run the algorithm for many iterations, the script may run for quite a while.

To give you an idea, I timed how much time my example took (I had a 535 * 535 image and I wanted 8 colors and ran the algorithm for 10 iterations):

time ./main.m ~/Documents/originalPicture.jpg ~/Documents/newPicture.jpg 8 10

gives me back approximately 9 minutes and 45 seconds.

I hope this post was a pretty good and fun introduction to the K-means algorithm. Do not hesitate to reach me if you encounter difficulties running or understanding the code.