#!/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
from sklearn.cluster import KMeans
from PIL import Image, ImageCms

#Open image and discard alpha channel which makes wheel round rather than square
im = Image.open('P355F1.jpg').convert('RGB')

# Convert to Lab colourspace
srgb_p = ImageCms.createProfile("sRGB")
lab_p  = ImageCms.createProfile("LAB")

rgb2lab = ImageCms.buildTransformFromOpenProfiles(srgb_p, lab_p, "RGB", "LAB")
temp = ImageCms.applyTransform(im, rgb2lab)
Y=np.asarray(temp)

#%%
# At this point, Y is now an 3-d array that we can work with.

temp1=Y.shape
nrows=temp1[0]
ncols=temp1[1]

Y1=Y[:,:,:2]
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)