import gradio as gr
import numpy as np
from sklearn.cluster import KMeans

with gr.Blocks() as demo:
    image = gr.Image("vulture.webp")
    n_colors = gr.Slider(1, 16, 4)
    button = gr.Button()
    plot = gr.Plot()

    @button.click(inputs=image, outputs=gr.JSON())
    def get_palette(image_array: np.ndarray):
        w, h, d = image_array.shape
        pixels = image_array.reshape(w * h, d)
        model = KMeans(n_clusters=4).fit(pixels)

        return model.cluster_centers_


if __name__ == "__main__":
    demo.launch()