Butterfly Image Segmentation Using K-means Clustering

In digital image processing and computer vision, image segmentation is the process of partitioning a digital image into multiple segments (sets of pixels, also known as image objects). (Wikipedia)

One commonly used image segmentation technique is K-means clustering. K-means clustering is a machine learning technique that separates an image into segments by clustering or grouping together data points that have similar traits.

In a nutshell, here's what K-means clustering is all about. In the context of computer vision, where we can represent and plot an image on 3d space, where each of the axes corresponds to one of the three color channels (RGB). 

1. Choose k random center points. Here we simply opt for 3 center points. 
   
2. We then assign every data point in the cluster based on the nearest center point as shown below 
   
3. Take the mean of all the values in each cluster. These new values then become the new center points. 

   

   
   
4. Repeat until convergence. 

Now that we have a brief idea of how the algorithm works, let's start with an image of a pretty butterfly.

We import it using the following code:

In order to prepare the data for k-means clustering, we need to reshape the data a bit and convert the values to float type using the following code:

Finally we implement k-means clustering algorithm. 

The segmented image looks like the following, where we only see 3 colors.
 

If we want to see each cluster in the image, we mask each of the clusters and set the color to green. 

Cluster 1: Orange/yellow part of the butterfly's wings as well as the white and purple flowers

Cluster 2: Dark parts of the butterfly's wings as well as dark part of the background

Cluster 3: Leaves

One can imagine how easy it is to play with the number of clusters. Increasing the clusters would allow us to compartment more of the details in the image. This is part of Udacity's Nanodegree program in Computer Vision.