# Importing Necessary Libraries
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from skimage import data
from skimage.color import rgb2gray

img = np.random.randint(10, size=(5, 6))
img

plt.figure(figsize=(3, 4))
plt.imshow(img, cmap=plt.cm.gray)

seg_img = img > 5
type(seg_img), seg_img

# plot each image ...
# ... side by side
fig = plt.figure(figsize=(6, 4))
fig.add_subplot(1, 2, 1)  # subplot one
plt.imshow(img, cmap=plt.cm.gray, vmin=0, vmax=9)

fig.add_subplot(1, 2, 2)  # subplot two
# my data is OK to use gray colormap (0:black, 1:white)
plt.imshow(seg_img, cmap=plt.cm.gray, vmin=0, vmax=1)  # use appropriate colormap here

integer_array = np.array([1, 2, 3, 7, 8, 9, 10])
X = integer_array.reshape(-1, 1)
# Create an instance of the KMeans class
kmeans = KMeans(n_clusters=2)
# Fit the K-means model to the data
kmeans.fit(X)
# Get the labels assigned by the K-means algorithm
labels = kmeans.labels_
# Get the cluster cluster_centers
cluster_centers = kmeans.cluster_centers_

print(labels)
print(cluster_centers)

kmeans = KMeans(n_clusters=2)
kmeans.fit(img)
labels = kmeans.labels_
cluster_centers = kmeans.cluster_centers_
print(labels)
print(cluster_centers)

# img.flatten('F')
X = img.reshape(-1,1)
kmeans.fit(X)
labels = kmeans.labels_
cluster_centers = kmeans.cluster_centers_
print(labels)
print(cluster_centers)

seg_img_labels = labels.reshape(img.shape)
seg_img_labels.shape

fig = plt.figure(figsize=(6, 4))
plt.imshow(seg_img_labels, cmap=plt.cm.gray)  

seg_img_labels

# seg_img_labels is the labels that was reshaped
seg_img_centers = np.zeros_like(img)
seg_img_centers[seg_img_labels == 0] = cluster_centers[0]
seg_img_centers[seg_img_labels == 1] = cluster_centers[1]
cluster_centers, seg_img_centers
# type(cluster_centers[0,0]), type(seg_img_centers[0,0])

fig = plt.figure(figsize=(6, 4))
fig.add_subplot(1, 2, 1)  # subplot one
plt.imshow(seg_img_centers, cmap=plt.cm.gray, vmin=0, vmax=9)  # use appropriate colormap here

fig.add_subplot(1, 2, 2)  # subplot two
plt.imshow(seg_img_labels, cmap=plt.cm.gray)  # use appropriate colormap here

# !pip install scikit-image

# Sample Image of scikit-image package
coffee = data.coffee()
coffee.shape, type(coffee[0, 0, 0])

# Setting the plot size to 15,15
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)

# Displaying the sample image
plt.imshow(coffee)

# Converting RGB image to Monochrome
gray_coffee = rgb2gray(coffee)
plt.subplot(1, 2, 2)

# Displaying the sample image - Monochrome
# Format
plt.axis("off")
plt.imshow(gray_coffee, cmap="gray")

X = gray_coffee.reshape(-1,1)
kmeans = KMeans(n_clusters=4)
kmeans.fit(X)
labels = kmeans.labels_
seg_img_labels = labels.reshape(gray_coffee.shape)

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.axis("off")
plt.imshow(gray_coffee, cmap="gray")

plt.subplot(1, 2, 2)
plt.axis("off")
plt.imshow(seg_img_labels, cmap="gray", vmin=0, vmax=3)

kmeans = KMeans(n_clusters=4)
kmeans.fit(X)
labels = kmeans.labels_
cluster_centers = kmeans.cluster_centers_

## New approach for replacing each pixel 
# with corresponding cluster center

seg_img_centers = cluster_centers[labels]
seg_img_centers = seg_img_centers.reshape(gray_coffee.shape)

plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.axis("off")
plt.imshow(gray_coffee, cmap="gray")

plt.subplot(1, 2, 2)
plt.axis("off")
plt.imshow(seg_img_centers, cmap="gray", vmin=0, vmax=1)