Using KMeans Clustering for Image Compression

KMeans_ImageCompression.py

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# Loading the image
org_image = Image.open('Kheerganga.jpeg','r')
pixels = list(org_image.getdata())
final_pixels = []
for i in range(0,len(pixels)):
    final_pixels.append( np.array(pixels[i]) )
final_pixels = np.array(final_pixels)

# importing scikit-learn for KMeans
from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters = 30, init = 'k-means++', n_init= 10, max_iter = 20)
y_kmeans = kmeans.fit_predict(final_pixels)

# reshaping it for Column*Rows
ykmeans = np.reshape(y_kmeans,(org_image.size[1],org_image.size[0]))
centroids = kmeans.cluster_centers_

# Plotting the compressed image clusters
X = np.arange(0,org_image.size[0])
for y in range(0,org_image.size[1]):
    plt.scatter(np.full((1,org_image.size[0]),y)[0],X ,c = kmeans.cluster_centers_[ykmeans[y][X]]/255)
plt.show()