Update app.py

app.py

@@ -4,11 +4,9 @@ import numpy as np
 import cv2
 from huggingface_hub import hf_hub_download
 import time
+import os
 
-# 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)
-
+# Registrar las funciones personalizadas
 from tensorflow.keras.saving import register_keras_serializable
 
 @register_keras_serializable()
@@ -24,14 +22,51 @@ def inverse_fourier_transform(x):
     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.")
+
+# 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 una sola vez
+model = tf.keras.models.load_model(model_path)
+
+# Crear versión optimizada para inferencia
+@tf.function(jit_compile=True)
+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)
+    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)
+    # Volver a escalarla al tamaño original
     degraded_img = cv2.resize(small_img, (w, h), interpolation=cv2.INTER_NEAREST)
 
     return degraded_img
@@ -39,7 +74,15 @@ def degrade_image(image, downscale_factor=4):
 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
+    # 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
@@ -49,55 +92,85 @@ def preprocess_image(image, std_dev=0.1, downscale_factor=4, target_size=(256, 2
     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]
+    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
     
-    # 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
+    # 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 utilizando la función optimizada
+        reconstructed = predict_optimized(noisy_image_input).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
+
+# Ejemplo para precalentamiento (warm-up) del modelo
+try:
+    dummy_input = np.zeros((1, 256, 256, 3), dtype=np.float32)
+    _ = predict_optimized(dummy_input)
+    print("Modelo precalentado con éxito")
+except Exception as e:
+    print(f"Error en precalentamiento: {e}")
+
+# Crear interfaz en Gradio
 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")],
+    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."
+    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
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=False)