np.random.seed(42)
img = cv.imread('images/flower.png')
gray_img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
plt.imshow(gray_img, cmap="gray")
plt.title("original image")

fig, axs = plt.subplots(1, 3, figsize=(15,15))

for i, k in enumerate([2,4,6]):
    centers, index = Kmeans(gray_img.reshape(-1,1), k=k)
    reconstructed_img = centers[index].reshape(gray_img.copy().shape)
    _ = axs[i].imshow(reconstructed_img/255)
    _ = axs[i].set_title("reconstructed_img k="+str(k))

img = cv.imread('images/flower.png')
plt.imshow(img)
plt.title("original image")

fig, axs = plt.subplots(1, 3, figsize=(15,15))

for i,k in enumerate([2,4,6]):
    centers, index = Kmeans(img.reshape(-1,3), k=k)
    reconstructed_img = centers[index].reshape(img.copy().shape)
    _ = axs[i].imshow(reconstructed_img/255)
    _ = axs[i].set_title("reconstructed_img k="+str(k))

img = cv.imread('images/flower.png')
gray_img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

plt.imshow(gray_img, cmap="gray")
plt.title("original image")

xx, yy = np.meshgrid(np.arange(gray_img.shape[1]), np.arange(gray_img.shape[0]))
xx = xx/np.max(xx)
yy = yy/np.max(yy)
new_img = gray_img/np.max(gray_img)
data = np.stack((xx, yy, new_img), axis=2)

plt.imshow(cv.cvtColor(data.astype(np.float32), cv.COLOR_BGR2RGB))
plt.title("original image")

fig, axs = plt.subplots(1, 3, figsize=(15,15))

for i,k in enumerate([2,4,6]):
    centers, index = Kmeans(data.reshape(-1,3), k=k)
    reconstructed_img = centers[index].reshape(data.copy().shape)[:,:,:]
    _ = axs[i].imshow(reconstructed_img)
    _ = axs[i].set_title("reconstructed_img k="+str(k))