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

# --- Cargar el modelo completo ---
model = tf.keras.models.load_model("autoencoder_complete_model_Fourier.keras")

# --- Transformaciones necesarias ---
def preprocess_image(img):
    """ Prepara la imagen para el modelo (escala 0-1 y redimensiona a 256x256). """
    img = img.resize((256, 256))  # Asegurar tamaño correcto
    img_array = np.array(img) / 255.0  # Normalizar entre 0 y 1
    img_array = np.expand_dims(img_array, axis=0)  # Añadir batch dimension
    return img_array

def denoise_image(img):
    """ Pasa la imagen ruidosa por el autoencoder y devuelve la imagen reconstruida. """
    img_array = preprocess_image(img)

    # Pasar la imagen por el autoencoder
    denoised_img = model.predict(img_array)

    # Postprocesamiento
    denoised_img = np.squeeze(denoised_img)  # Remover batch dimension
    denoised_img = (denoised_img * 255).astype(np.uint8)  # Convertir a formato uint8
    return denoised_img

# --- Interfaz Gradio ---
gr.Interface(fn=denoise_image,
             inputs=gr.Image(type="pil"),  # Imagen de entrada
             outputs=gr.Image(type="numpy"),  # Imagen de salida
             examples=["noisy1.jpg", "noisy2.jpg"]  # Ejemplos opcionales
).launch()