import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Carga el modelo entrenado
model = tf.keras.models.load_model("quickdraw_model.keras")
etiquetas = ['apple', 'banana', 'bed', 'carrot', 'laptop']

def preprocesar_imagen(image):
    # Convierte la imagen a escala de grises, 28x28 y binariza
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert('L')
    image = image.resize((28, 28), Image.NEAREST)
    arr = np.array(image) / 255.0
    arr_bin = (arr < 0.5).astype(np.float32)
    arr_bin_4d = arr_bin.reshape(1, 28, 28, 1)
    return arr_bin_4d

def predict(image):
    x = preprocesar_imagen(image)
    preds = model.predict(x)
    class_idx = np.argmax(preds)
    confidences = {etiquetas[i]: float(preds[0][i]) for i in range(len(etiquetas))}
    return confidences

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Dibuja o sube una imagen (fondo blanco, trazo negro)"),
    outputs=gr.Label(label="Predicción"),
    title="QuickDraw - Clasificador de Dibujos",
    description="Dibuja o sube una imagen y te dirá si es una apple, banana, bed, carrot o laptop."
)

iface.launch()