# -*- coding: utf-8 -*-
"""Copy of extract_colors_from_image.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Zx45R30-L2sIBh8VU_Fnbv9G6v4445TD
"""

# Commented out IPython magic to ensure Python compatibility.


from sklearn.cluster import KMeans
from collections import Counter
import numpy as np
import cv2

import gradio as gr


def get_image(pil_image):
    #image = cv2.imread(image_path)
    nimg = np.array(pil_image)
    image = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image

def get_labels(rimg):
  clf = KMeans(n_clusters = 5)
  labels = clf.fit_predict(rimg)
  return labels ,  clf

def RGB2HEX(color):
    return "#{:02x}{:02x}{:02x}".format(int(color[0]), int(color[1]), int(color[2]))

def get_colors(pimg):
  img = get_image(pimg)
  reshaped_img = img.reshape(img.shape[0]*img.shape[1], img.shape[2])
  labels, clf = get_labels(reshaped_img)
  counts = Counter(labels)
  center_colors = clf.cluster_centers_
  # We get ordered colors by iterating through the keys
  ordered_colors = [center_colors[i] for i in counts.keys()]
  hex_colors = [RGB2HEX(ordered_colors[i]) for i in counts.keys()]
  #rgb_colors = [ordered_colors[i] for i in counts.keys()]
  return hex_colors

demo = gr.Blocks()

with demo:
    gr.Markdown(
        """
    # Extract Colors from an image using KMeans clustering
    """
    )
    inputs = [gr.Image(type="pil", label="Image to extract colors from")]
    with gr.Row():
        outputs = [gr.ColorPicker(label="color 1"), gr.ColorPicker(label="color 2"),gr.ColorPicker(label="color 3"),gr.ColorPicker(label="color 4"),gr.ColorPicker(label="color 5")]

    btn = gr.Button("Extract colors")
    btn.click(fn=get_colors, inputs=inputs, outputs=outputs)
demo.queue()
demo.launch()