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

# Cargar modelo
model = tf.keras.models.load_model("quickdraw_model.keras")
etiquetas = ['apple', 'banana', 'bed', 'carrot', 'laptop']  # Ajusta según tu modelo

def preprocesar_imagen(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if image.mode in ('RGBA', 'LA'):
        image = image.convert('RGB')
    image = image.convert('L')  # escala de grises
    image = image.resize((28, 28), Image.NEAREST)
    arr = np.array(image) / 255.0
    
    # Umbral bajo: todo menor o igual a 0.2 será negro (0), lo demás blanco (1)
    arr_bin = np.where(arr <= 0.2, 0.0, 1.0).astype(np.float32)
    
    arr_bin_4d = arr_bin.reshape(1, 28, 28, 1)
    
    return arr_bin_4d, arr_bin

def predict(image):
    try:
        arr, img_preprocesada = preprocesar_imagen(image)
        preds = model.predict(arr)
        class_idx = np.argmax(preds)
        return {etiquetas[class_idx]: float(preds[0][class_idx])}, img_preprocesada
    except Exception as e:
        return f"Error: {str(e)}", None

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Dibuja o sube una imagen (fondo blanco, trazo negro)"),
    outputs=[
        gr.Label(num_top_classes=1, label="Predicción"),
        gr.Image(label="Imagen preprocesada (0-1 en escala de grises)")
    ],
    title="QuickDraw API",
    description="API para reconocer dibujos estilo QuickDraw. Muestra la imagen preprocesada en escala 0-1."
)

iface.launch()