#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Mar 21 16:31:10 2021

Example of image segmentation based on k-means
clustering of data in color space.

@author: doug
"""
import numpy as np
import matplotlib.pyplot as plt
import scipy.io                    # To read Mat Files
from sklearn.cluster import KMeans


X=scipy.io.loadmat('ImageAB.mat')
Y1=X['Y1']
nrows=int(X['nrows'])
ncols=int(X['ncols'])

Y2=np.reshape(Y1,(nrows*ncols,2))

kmeans=KMeans(init="random",n_clusters=3,n_init=10)
kmeans.fit(Y2)

disterr=kmeans.inertia_
Clusters=kmeans.cluster_centers_
Index=kmeans.labels_

Z=np.reshape(Index,(nrows,ncols))
plt.imshow(Z)