import numpy as np
import cv2
from sklearn.cluster import MeanShift, estimate_bandwidth

#Loading original image
original_img = cv2.imread('Swimming_Pool.jpg')

#Shape of original image
original_shape = original_img.shape

#Converting image into array of dimension [nb of pixels in original_img, 3] feature array based on r g b intensities
flat_img=np.reshape(original_img, [-1, 3])

#Estimate bandwidth for meanshift algorithm
bandwidth = estimate_bandwidth(flat_img, quantile=0.1, n_samples=20)
ms = MeanShift(bandwidth = bandwidth, bin_seeding=True)

#Performing meanshift on flatImg
ms.fit(flat_img)

#(r,g,b) vectors corresponding to the different clusters after meanshift
labels=ms.labels_

#Remaining colors after
cluster_centers = ms.cluster_centers_

#Finding and diplaying the number of clusters
labels_unique = np.unique(labels)
n_clusters_ = len(labels_unique)
print("number of estimated clusters : %d" % n_clusters_)

#Displaying segmented image
segmented_img = np.reshape(labels, original_shape[:2])
cv2.imshow('Image',segmented_img)
cv2.waitKey(0)
cv2.destroyAllWindows()