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

def flatten_colors(input_img, num_colors):
    if input_img is None:
        return None
    
    # Convert PIL image to numpy array
    img_array = np.array(input_img)
    original_shape = img_array.shape
    
    # Reshape to a list of pixels (R, G, B)
    pixels = img_array.reshape(-1, 3)
    
    # Apply K-Means to find dominant colors
    kmeans = KMeans(n_clusters=int(num_colors), random_state=42, n_init=10)
    labels = kmeans.fit_predict(pixels)
    colors = kmeans.cluster_centers_.astype('uint8')
    
    # Map each pixel to its corresponding cluster center (the flat color)
    flattened_pixels = colors[labels]
    
    # Reshape back to original image dimensions
    flattened_img = flattened_pixels.reshape(original_shape)
    
    return Image.fromarray(flattened_img)

# Define the Gradio Interface
demo = gr.Interface(
    fn=flatten_colors,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(minimum=2, maximum=20, value=8, step=1, label="Number of Colors")
    ],
    outputs=gr.Image(type="pil", label="Flat Color Result"),
    title="🎨 Image Color Flattener",
    description="Upload an image and reduce it to a flat color palette using K-Means clustering."
)

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