Spaces:
Runtime error
Runtime error
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
import cv2 | |
from huggingface_hub import hf_hub_download | |
import time | |
import os | |
# Registrar las funciones personalizadas | |
from tensorflow.keras.saving import 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) | |
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)) | |
# Configuración de GPU para TensorFlow | |
physical_devices = tf.config.list_physical_devices('GPU') | |
if physical_devices: | |
print("GPU disponible. Configurando...") | |
try: | |
# Permitir crecimiento de memoria según sea necesario | |
for gpu in physical_devices: | |
tf.config.experimental.set_memory_growth(gpu, True) | |
print("Configuración de GPU completada") | |
except Exception as e: | |
print(f"Error en configuración de GPU: {e}") | |
else: | |
print("No se detectó GPU. El procesamiento será más lento.") | |
# IMPORTANTE: Habilitar deserialización insegura para capas Lambda | |
try: | |
# Para TensorFlow 2.11+ | |
if hasattr(tf.keras.config, 'enable_unsafe_deserialization'): | |
tf.keras.config.enable_unsafe_deserialization() | |
print("Deserialización insegura habilitada mediante tf.keras.config") | |
# Para versiones antiguas, intentar Keras directamente | |
elif hasattr(tf.keras.utils, 'enable_unsafe_deserialization'): | |
tf.keras.utils.enable_unsafe_deserialization() | |
print("Deserialización insegura habilitada mediante tf.keras.utils") | |
else: | |
print("No se pudo habilitar la deserialización insegura automáticamente") | |
except Exception as e: | |
print(f"Error al configurar deserialización: {e}") | |
# Descargar modelo desde Hugging Face (con caché) | |
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface") | |
model_path = hf_hub_download( | |
repo_id="Bmo411/DenoisingAutoencoder", | |
filename="autoencoder_complete_model_Fourier.keras", | |
cache_dir=cache_dir | |
) | |
print(f"Modelo cargado desde: {model_path}") | |
# Cargar el modelo con desactivación de safety | |
try: | |
# Intentar cargar con custom_objects para las funciones personalizadas | |
model = tf.keras.models.load_model( | |
model_path, | |
custom_objects={ | |
'fourier_transform': fourier_transform, | |
'inverse_fourier_transform': inverse_fourier_transform | |
}, | |
compile=False # No compilar el modelo para inferencia más rápida | |
) | |
print("Modelo cargado correctamente") | |
except Exception as e: | |
print(f"Error al cargar el modelo: {e}") | |
# Alternativa fallback para versiones más recientes de TF | |
try: | |
options = tf.saved_model.LoadOptions( | |
experimental_io_device='/job:localhost' | |
) | |
model = tf.keras.models.load_model( | |
model_path, | |
custom_objects={ | |
'fourier_transform': fourier_transform, | |
'inverse_fourier_transform': inverse_fourier_transform | |
}, | |
compile=False, | |
options=options | |
) | |
print("Modelo cargado con opciones alternativas") | |
except Exception as e2: | |
print(f"Error al cargar el modelo con opciones alternativas: {e2}") | |
raise Exception("No se pudo cargar el modelo") | |
# Crear versión optimizada para inferencia | |
def predict_optimized(input_tensor): | |
return model(input_tensor, training=False) | |
# Funciones de preprocesamiento optimizadas | |
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] | |
# Verificar tamaño mínimo | |
if h < downscale_factor*4 or w < downscale_factor*4: | |
return image # Evitar downscaling excesivo en imágenes pequeñas | |
# 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 | |
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.""" | |
# Verificar si la imagen es a color o en escala de grises | |
if len(image.shape) == 2: | |
# Convertir a RGB si es escala de grises | |
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) | |
elif image.shape[2] == 4: | |
# Convertir de RGBA a RGB si tiene canal alpha | |
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) | |
# Reducir calidad | |
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) | |
noisy_img = resized_img + noise | |
noisy_img = np.clip(noisy_img, 0, 1) | |
return noisy_img | |
# Variable para medir el tiempo total de la primera ejecución | |
first_run = True | |
def Denoiser(imagen): | |
"""Aplica el modelo autoencoder para eliminar ruido de la imagen.""" | |
global first_run | |
# Verificar que la imagen no sea None | |
if imagen is None: | |
return None, None | |
# Convertir imagen de entrada a array NumPy | |
imagen = np.array(imagen) | |
# Verificar dimensiones de la imagen | |
if len(imagen.shape) < 2: | |
return None, None | |
try: | |
# Preprocesar imagen (degradarla y agregar ruido) | |
noisy_image = preprocess_image(imagen) | |
# Expandir dimensiones para el formato del modelo | |
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 | |
# Usar llamada directa en lugar de predict() para más velocidad | |
reconstructed = model(noisy_image_input, training=False).numpy()[0] | |
prediction_time = time.time() - start_time | |
if first_run: | |
print(f"Primera ejecución: {prediction_time:.2f} segundos") | |
first_run = False | |
else: | |
print(f"Tiempo de predicción: {prediction_time:.2f} segundos") | |
# Asegurarse de que las imágenes estén en el rango [0, 255] | |
noisy_image = np.uint8(noisy_image * 255) | |
reconstructed = np.uint8(reconstructed * 255) | |
return noisy_image, reconstructed | |
except Exception as e: | |
print(f"Error en el procesamiento: {e}") | |
return None, None | |
# Crear interfaz en Gradio | |
demo = gr.Interface( | |
fn=Denoiser, | |
inputs=gr.Image(type="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.", | |
examples=[ | |
"https://raw.githubusercontent.com/gradio-app/gradio/main/demo/english_htr/images/Create%20a%20free%20Gradio%20account%20to%20access%20our%20most%20powerful%20features.jpeg" | |
], | |
cache_examples=True | |
) | |
# Lanzar la aplicación en Gradio | |
if __name__ == "__main__": | |
demo.launch(share=False) |