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

# Cargar el modelo DAE desde Hugging Face
modelo_path = hf_hub_download(repo_id="XTEP63/DAE", filename="denoiser.keras")
dae_model = tf.keras.models.load_model(modelo_path)

# Función para procesar la imagen con el DAE
def denoise_image(image):
    image = image.resize((128, 128))  # Ajusta según el tamaño de las imagenes
    image = np.array(image) / 255.0  # Normalizar
    image = np.expand_dims(image, axis=0)  
    output = dae_model.predict(image)
    output = np.squeeze(output, axis=0)  
    output = (output * 255).astype(np.uint8)  # Volver a escala 0-255
    return Image.fromarray(output)

# Interfaz 
iface_dae = gr.Interface(
    fn=denoise_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Denoising Autoencoder",
    description="Sube una imagen con ruido y el modelo la limpiará."
)

iface_dae.launch()