app.py

import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from huggingface_hub import hf_hub_download
import time

# Descargar modelo desde Hugging Face
model_path = hf_hub_download(repo_id="Bmo411/DenoisingAutoencoder", filename="autoencoder_complete_model_Fourier.keras")
model = tf.keras.models.load_model(model_path)

from tensorflow.keras.saving import register_keras_serializable

@register_keras_serializable()
def fourier_transform(x):
    fourier = tf.signal.fft2d(tf.cast(x, tf.complex64))
    fourier = tf.complex(tf.math.real(fourier), tf.math.imag(fourier))
    fourier = tf.abs(fourier)
    return tf.concat([tf.math.real(fourier), tf.math.imag(fourier)], axis=-1)

@register_keras_serializable()
def inverse_fourier_transform(x):
    real_part, imag_part = tf.split(x, num_or_size_splits=2, axis=-1)
    complex_fourier = tf.complex(real_part, imag_part)
    return tf.abs(tf.signal.ifft2d(complex_fourier))

def degrade_image(image, downscale_factor=4):
    """Reduce la calidad de la imagen reduciendo su tamaño y volviéndola a escalar."""
    h, w = image.shape[:2]

    # Reducir tamaño (forzando pérdida de calidad)
    small_img = cv2.resize(image, (w // downscale_factor, h // downscale_factor), interpolation=cv2.INTER_AREA)

    # Volver a escalarla al tamaño original SIN mejorar calidad (interpolación brusca)
    degraded_img = cv2.resize(small_img, (w, h), interpolation=cv2.INTER_NEAREST)

    return degraded_img

def preprocess_image(image, std_dev=0.1, downscale_factor=4, target_size=(256, 256)):
    """Recibe una imagen en numpy array, la degrada en calidad, le agrega ruido y la normaliza."""
    
    # Reducir calidad primero
    degraded_img = degrade_image(image, downscale_factor)

    # Redimensionar a tamaño esperado por el modelo
    resized_img = cv2.resize(degraded_img, target_size, interpolation=cv2.INTER_AREA)
    
    # Normalizar
    resized_img = resized_img.astype(np.float32) / 255.0

    # Agregar ruido gaussiano
    noise = np.random.normal(0, std_dev, resized_img.shape)  # Media=0, desviación estándar=std_dev
    noisy_img = resized_img + noise  # Sumar el ruido a la imagen
    noisy_img = np.clip(noisy_img, 0, 1)  # Asegurar valores en el rango [0,1]

    return noisy_img

def Denoiser(imagen):
    """Aplica el modelo autoencoder para eliminar ruido de la imagen."""
    
    # Convertir imagen de entrada a array NumPy y asegurarse de que está en formato correcto
    imagen = np.array(imagen)  # Convertir de PIL a NumPy

    # Preprocesar imagen (degradarla y agregar ruido)
    noisy_image = preprocess_image(imagen)

    # Expandir dimensiones para que tenga el formato correcto para el modelo (batch_size, h, w, c)
    noisy_image_input = np.expand_dims(noisy_image, axis=0)

    # Medir el tiempo de la predicción
    start_time = time.time()

    # Predecir con el autoencoder
    reconstructed = model.predict(noisy_image_input)[0]  # Quitar batch_size

    prediction_time = time.time() - start_time
    print(f"Tiempo de predicción: {prediction_time} segundos")

    # Asegurarse de que la imagen restaurada esté en el rango [0, 255] para la visualización
    noisy_image = np.uint8(noisy_image * 255)
    reconstructed = np.uint8(reconstructed * 255)

    return noisy_image, reconstructed

# Verificar si TensorFlow está utilizando la GPU
physical_devices = tf.config.list_physical_devices('GPU')
if physical_devices:
    print("TensorFlow está utilizando la GPU")
else:
    print("TensorFlow no está utilizando la GPU")

# Crear interfaz en Gradio con dos salidas de imagen
demo = gr.Interface(
    fn=Denoiser, 
    inputs=gr.Image(type="numpy"),  # Para que Gradio pase la imagen como un array de NumPy
    outputs=[gr.Image(type="numpy", label="Imagen con Ruido"), 
             gr.Image(type="numpy", label="Imagen Restaurada")],
    title="Autoencoder para Denoising",
    description="Este modelo de autoencoder reduce el ruido en imágenes. Sube una imagen y el modelo generará una versión restaurada."
)

# Lanzar la aplicación en Gradio
demo.launch()