import gradio as gr
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
import os

# Definir rutas
code_path = '/app'  # Usando ruta del contenedor Docker
haarcascade_path = 'haarcascade_frontalface_default.xml'  # Usamos ruta relativa

# Rutas de los modelos
modelo_mejor_path = os.path.join(code_path, 'modelo_estimacion_edad_mejor.h5')
modelo_final_path = os.path.join(code_path, 'modelo_estimacion_edad_final.h5')

# Imprimir la ruta para verificar
print(f"Ruta al archivo Haar Cascade: {haarcascade_path}")

# Verificar si el archivo Haar Cascade existe
if not os.path.exists(haarcascade_path):
    print(f'No se encontró el archivo Haar Cascade en {haarcascade_path}. Asegúrate de subirlo correctamente en tu Space.')
    exit()

# Cargar el clasificador Haar Cascade para detección de rostros
face_cascade = cv2.CascadeClassifier(haarcascade_path)

# Verificar si Haar Cascade se cargó correctamente
if face_cascade.empty():
    print("Error al cargar el clasificador Haar Cascade.")
    exit()

# Cargar ambos modelos
modelo_mejor = load_model(modelo_mejor_path)
modelo_final = load_model(modelo_final_path)

print('Modelos cargados exitosamente.')

# Función para procesar la imagen y predecir la edad
def estimar_edad(imagen, modelo_seleccionado="mejor"):
    # Seleccionar el modelo
    if modelo_seleccionado == "mejor":
        modelo = modelo_mejor
    else:
        modelo = modelo_final

    # Convertir la imagen a escala de grises para la detección de rostros
    gray = cv2.cvtColor(imagen, cv2.COLOR_BGR2GRAY)

    # Detectar rostros en la imagen
    rostros = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    edades = []

    for (x, y, w, h) in rostros:
        # Extraer la región del rostro
        rostro = imagen[y:y+h, x:x+w]

        # Preprocesar la imagen del rostro
        rostro_rgb = cv2.cvtColor(rostro, cv2.COLOR_BGR2RGB)
        rostro_resized = cv2.resize(rostro_rgb, (128, 128))
        rostro_preprocessed = preprocess_input(rostro_resized)
        rostro_expanded = np.expand_dims(rostro_preprocessed, axis=0)

        # Realizar la predicción de edad
        edad_predicha = modelo.predict(rostro_expanded)[0][0]
        edad_predicha = round(edad_predicha, 2)

        edades.append(edad_predicha)

    if len(edades) > 0:
        return f"Edad estimada: {round(np.mean(edades), 2)} años"
    else:
        return "No se detectaron rostros en la imagen."

# Crear la interfaz con Gradio (desactivando flagging)
iface = gr.Interface(
    fn=estimar_edad, 
    inputs=[
        gr.Image(type="numpy", label="Sube una imagen"),
        gr.Radio(choices=["mejor", "final"], label="Selecciona el modelo")  # Elegir modelo
    ], 
    outputs="text",
    flagging="never"  # Desactivar flagging para evitar la creación del directorio 'flagged'
)

# Lanzar la interfaz de Gradio
iface.launch()