362. Image Segmentation using K-means Clustering

Clustering

Depending on your data and objective, you may not even have to train a deep-learning model for image segmentation. Here is one way to apply segmentation using Kmeans clustering.

Implementation

from sklearn.cluster import KMeans
from matplotlib.image import imread
import matplotlib.pyplot as plt
import numpy as np

#Read Image
image = imread("/content/drive/MyDrive/co.jpg")

#Flatten Channel-Wise
X = image.reshape(-1,3)

#Assign Number of Clusters
k = 2

#Fit
kmeans = KMeans(n_clusters=k).fit(X)

#Display Results
segmented_img = kmeans.cluster_centers_[kmeans.labels_]
segmented_img = segmented_img.reshape(image.shape)
plt.imshow((segmented_img * 255).astype(np.uint8))