# Importing the Gradio library for creating interactive web interfaces
import gradio as gr

# Importing necessary libraries for image processing
import numpy as np
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

# Function to perform k-means color quantization on an image
def knn_color_quantization(image, n_colors):
    # Reshaping the image into a 2D array of pixels
    pixels = np.reshape(image, (image.shape[0] * image.shape[1], image.shape[2]))

    # Performing k-means clustering to find the dominant colors
    kmeans = KMeans(n_clusters=n_colors, random_state=42)
    kmeans.fit(pixels)

    # Getting the cluster centers as the representative colors
    kmeans_centers = kmeans.cluster_centers_

    # Fitting a nearest neighbors model to find the closest color for each pixel
    knn = NearestNeighbors(n_neighbors=1)
    knn.fit(kmeans_centers)

    # Finding the closest color for each pixel and reconstructing the quantized image
    distances, indices = knn.kneighbors(pixels)
    quantized_pixels = kmeans_centers[indices.flatten()]
    quantized_image = np.reshape(quantized_pixels, image.shape)

    # Scaling the image values to the range [0, 1]
    quantized_image = quantized_image / 255.0

    return quantized_image

# Function to create the Gradio interface for color quantization
def color_quantization_app(image, num_colors):
    # Calling the knn_color_quantization function to quantize the image
    quantized_image = knn_color_quantization(image, num_colors)
    # Converting the pixel values back to the range [0, 255] and casting to uint8
    quantized_image = (quantized_image * 255).astype(np.uint8)
    return quantized_image
    
dark_minimalist = gr.Theme.from_hub("Taithrah/Minimal")

# Creating a Gradio interface object
iface = gr.Interface(
    # Specifying the function to be executed
    fn=color_quantization_app,
    # Defining the inputs for the interface (an uploaded image and a slider for number of colors)
    inputs=[gr.Image(label="Uploaded Image", sources='upload', type="numpy"), gr.Slider(minimum=2,maximum=300,value=8, label='Number of Colors',interactive=True)],
    # Defining the output for the interface (quantized image)
    outputs=gr.Image(label="Quantized Image"),
    # Disabling live updates
    live=False,
    # Disabling user flagging of outputs
    allow_flagging="never",
    # Applying a custom theme retrieved from Gradio Hub
    theme=dark_minimalist,
    # Adding custom CSS to hide the footer of the interface
    css="""
    footer {
        visibility: hidden;
    }
    """
)

# Launching the Gradio interface
iface.launch()