Image Segmentation using K-means clustering algorithm | Python

· 3 min read
Image Segmentation using K-means clustering algorithm | Python

Ina previous article, we saw how to implement the K-means algorithm from scratch in python. We delved deep into the working of the algorithm and discussed some possible practical applications. In this tutorial, we are going to see one such application at work. In this tutorial, we will see how we can use K-means clustering to separate an image into segments based on its pixel values.

If you are new to machine learning or K-means, you can read the original article here.

The complete code used in this article can be found here.

Since we have discussed all the nitty-gritty details in the original article, I will keep this one fairly short and simple.

This tutorial will be split into 3 short portions.

· Preprocessing step for the image files.

· Working of the Algorithm.

· Results.

Let’s begin!!

Loading our image will be our first step.

We will use the following image downloaded from

Image of a red car
image credits: Josh Rinard (

We cannot deal with images the same way as we dealt with normal data points, this is because colored images are a 3-dimensional matrix (The third dimension is the RGB color channels and does not add to the perspective but this is important for our mathematics to work).

RGB peresentation of an image
Dimensional representation of an image

Each channel consists of pixels containing information regarding the channel color (Red, Green, or Blue)

We need to convert this into 2-dimensional data.

We now have the flattened data in a data frame. It is time to write the algorithm. The Algorithm will remain the same as the original one before, for an in-depth look into K-means clustering, read the original article here.

The algorithm is run for 5 clusters; as denoted by the variable ‘k’. Instead of waiting for the centroids to reach a stable value, this time we have set a threshold to ensure we have optimal centroids.

This is done because, for image data, the algorithm takes a very long time to find the perfect centroid, so we give it a small margin i.e. we say that if the difference between the new and the old centroids is within a certain range, quit the loop and keep the last found centroids.

It is important to note at this point that the algorithm we have written above is quite slow in its working since speed and efficiency are not the goals of this tutorial, here we are only looking at the working of the K-means algorithm. One way to improve the speed of the algorithm is to use the “itertuples” function instead of the “iterrows” since tuples are at least ten times faster in traversal due to their static memory allocation. If you know any other ways to improve the efficiency of this algorithm, do comment down below.

After a little wait, the algorithm has finished its work, lets see what centroids we have in the end result.

centroids = centroids.to_numpy()
Centroid locations
Final centroids found

Now that we have the centroids, all we need to do is plot the result. The plotting is done in such a way that all the pixels in a particular cluster are overwritten by the centroid of that cluster.

Let's see what the final image looks like.

Segmented image
Final image


The final image has only 5 colors in total (due to 5 clusters), these 5 colors represent the major colors that were present in the original image.

The final image looks like something out of an Instagram filter. We can use the centroids and the clusters obtained to create multiple masks for the image and separate it into portions, e.g. we can see that the body of the car is red but split into 2 different tones, We can train the algorithm for a fewer number of clusters and then separate the red body of the car altogether since it will part of a single cluster.

Similarly, we can find multiple use cases for image segmentation in this way.