import gradio as gr
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.losses import MeanSquaredError
from PIL import Image

# Cargar modelo con custom_objects para mapear 'mse'
model = load_model("autoencoder.h5", custom_objects={"mse": MeanSquaredError()})

# Función de predicción
def detectar_anomalia(imagen):
    # Convertir y preprocesar la imagen
    imagen = imagen.convert("L").resize((64, 64))
    arr = np.array(imagen) / 255.0
    arr = arr.reshape((1, 64, 64, 1))
    
    # Obtener la reconstrucción
    reconstruido = model.predict(arr, verbose=0)  # verbose=0 para evitar logs innecesarios
    
    # Calcular el error de reconstrucción (MSE)
    error = np.mean((arr - reconstruido) ** 2)
    
    # Convertir las imágenes para visualización
    img_original = (arr.squeeze() * 255).astype(np.uint8)
    img_reconstruida = (reconstruido.squeeze() * 255).astype(np.uint8)
    
    # Convertir a objetos PIL para Gradio
    imagen_original = Image.fromarray(img_original)
    imagen_reconstruida = Image.fromarray(img_reconstruida)

    # Determinar si es anómala (usando un umbral)
    is_anomaly = error > 0.01  # Ajusta este umbral según tus necesidades
    return (imagen_original, 
            imagen_reconstruida, 
            f"Error de reconstrucción (MSE): {error:.6f}\n¿Es anómala?: {'Sí' if is_anomaly else 'No'}")

# Interfaz de Gradio
demo = gr.Interface(
    fn=detectar_anomalia,
    inputs=gr.Image(type="pil", label="Sube una imagen para analizar"),
    outputs=[
        gr.Image(label="Imagen Original"),
        gr.Image(label="Imagen Reconstruida"),
        gr.Textbox(label="Resultado de la Detección")
    ],
    title="Detección de Anomalías en Textiles con Autoencoder (Keras)",
    description="Este Space utiliza un autoencoder entrenado con Keras para detectar anomalías en imágenes de textiles industriales. Sube una imagen para comparar la original con su reconstrucción y determinar si contiene defectos (como manchas, texto o patrones irregulares) basándote en el error de reconstrucción (MSE). Un MSE alto indica una posible anomalía."
)

demo.launch()