import os
import cv2
import numpy as np
import imghdr
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from PIL import Image
plt.style.use("default")


class Segmentation:
    '''
    Image segmentation work
    '''

    def __init__(self, image, dom_colors):
        # since cv2 reads image in BGR mode , it is necessary to convert into RGB mode.
        self.image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        self.orginal_size = image.shape
        # flatenning the image in order to make it compatible for KMeans class
        self.pixel_array = self.image.reshape((-1, 3))
        self.dom_colors = dom_colors
        # Created an instance of KMeans
        self.km = KMeans(n_clusters=self.dom_colors)
        # created a model here
        self.km.fit(self.pixel_array)

    # this program helps to extract out the dominant colors from the image
    def dominant_colors(self):
        # taking out the centers
        self.centers = np.array(self.km.cluster_centers_, dtype='uint8')
        self.colors = []
        plt.figure(0, (4, 4))
        plt.axis("off")
        var = 1
        for current_center in self.centers:
            plt.subplot(1, self.dom_colors+1, var)
            self.colors.append(current_center)
            # created an array to store data of each dominant color
            color_array = np.zeros((100, 100, 3), dtype='uint8')
            color_array[:, :, :] = current_center
            plt.imshow(color_array)
            var += 1
        plt.show()

    # this fucntion draws the image with the given dominant colors
    def draw_image(self):
        self.centers = np.array(self.km.cluster_centers_, dtype='uint8')
        # here predict function gives label to each point , i.e the given point is nearer to which center
        pred = self.km.predict(self.pixel_array)
        # creaed an empty array to store data of image
        new_image = np.zeros(
            (self.image.shape[0]*self.image.shape[1], 3), dtype='uint8')
        for i in range(new_image.shape[0]):
            new_image[i] = self.centers[pred[i]]
        # new_image is reshaped into original size in order to get whole image together
        new_image = new_image.reshape(self.orginal_size)
        plt.axis("off")
        plt.title("No. of colors : "+str(self.dom_colors))
        plt.imshow(new_image)
        plt.savefig("a.png")
        plt.show()


def execute(addr):
    abspath = os.path.abspath(addr)
    try:
        im = Image.open(abspath)
        img = cv2.imread(abspath)
        IS = Segmentation(img, 2)
        IS.draw_image()
    except IOError:
        print("Error...Only JPEG,JPG,PNG formats are supported")
