Autoencoders / app.py
Bmo411's picture
Update app.py
9d2d3c1 verified
raw
history blame
7.9 kB
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
@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))
# 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
@tf.function
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)