Image Segmentation with K-Means Clustering in Python

Extract ‘superpixels’ of an Image using the clustering approach

Harsha S
4 min readFeb 9, 2023

Before we get into the Image Segmentation using K-Means clustering, let’s quickly brush upon the basics.

K-Means Clustering

The basic underlying idea behind any clustering algorithm is to partition a set of values into a specific number of cluster. In K-Means clustering, the k clusters are assigned with values that are nearest to the mean of a particular cluster. The steps are iterative with mean shifted in vector space until the convergence is achieved.

This was a short description of K-Means clustering to get us started, here is the more detailed version.

Image Segmentation

Image segmentation falls into a category of Image processing, which involves partitioning an image into multiple segments that have more similarity. This division into distinct regions is based on high similarity between pixels of each region and high contrast between regions.

The segmented image is more easier to analyze, especially in the case of medical images as a region of interest can be isolated from the background image.

Image Segmentation is the overall process of assigning a label to all of the pixels in an image such that pixels with the same label share certain characteristics. It’s useful in a variety of fields, including healthcare, digital image processing, traffic image analysis, and pattern recognition.

Now let’s implement the Image Segmentation via K-Means Clustering in Python using OpenCV library.

Import the necessary modules:

import cv2
import numpy as np
import matplotlib.pyplot as plt

Load the image and convert it to RGB color space

img = cv2.imread("/content/pexels-gabriel-frank-14224359.jpg")
img=cv2.cvtColor(img ,cv2.COLOR_BGR2RGB)
plt.figure(figsize=(13,10))
plt.imshow(img)

The Image before segmentation:

Reshape the image into a vector:

vectorized_img = img.reshape((-1,3))
vectorized_img.shape

In the next step, we convert the unit8 values to float type to satisfy the OpenCV requirements.

vectorized = np.float32(vectorized_img)

Define the number of clusters along with the criteria to finish the segmentation process.

k = 4
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 20, 0.1)
_, labels, centers = cv2.kmeans(vectorized, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)

To successfully perform the k-means clustering, we use the cv2.means function from OpenCV library.

It takes the following arguments respectively:

  • criteria — determines how long the process of segmentation works when trying to find the best cluster position.

There are three stop process criteria flags:

  1. cv.TERM_CRITERIA_EPS — stop the algorithm iteration if specified accuracy, epsilon, is reached.
  2. cv.TERM_CRITERIA_MAX_ITER — stop the algorithm after the specified number of iterations, max_iter.
  3. cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER — stop the iteration when any of the above condition is met.
  • input data — a vector of points with float values, here it is vectorized with three — dimensional points
  • number of clusters — a specification for the ‘k’ clusters. The higher number of clusters tend to produce less noise in the segmented image.
  • input/output labels — used to determine the initial positions of the clusters and get back the resulting clusters. It is set to None as we don’t want to initialize the cluster centers manually.
  • attempts — defines how many times the clusterization process will be launched from different cluster label initializations. It returns the best attempt with good compactness of the clusters.
  • flags — determine the type of cluster initialization
  1. cv2.KMEANS_RANDOM_CENTERS — random initialization
  2. cv2.KMEANS_PP_CENTERS — more sophisticated initialization
  3. cv2.KMEANS_USE_INITIAL_LABELS — user specified cluster centers

Apply the colors of centroids to all pixels that are relevant to these centroids. Afterwards, reshape the segmented image back to its original shape. Then, convert it to RGB color space

segmented = center[labels.flatten()].reshape(img.shape)

Display the original image and the segmented image (k=4) side by side.

plt.subplot(121)
plt.axis('off')
plt.title('Original Image')
plt.imshow(img)
plt.subplot(122)
plt.axis('off')
plt.title(f'Segmented Image when K = {k}')
plt.imshow(segmented)
plt.show()

Let’s now assign the value k = 6 and check the results

Let’s check the results for k = 10

We have segmented an image by using k-means clustering and obtained results for different values of k.

As it is evident, from the results the higher k value produces more clear segmented with more clusters, however the selection of k depends on the field and type of images we are dealing with.

Thanks for your time!

Follow for more articles on AI and ML!

BECOME a WRITER at MLearning.ai

--

--

Harsha S

NLP Engineer | I love to write about AI in beginner way