import os
import streamlit as st
import cv2
import numpy as np
import tensorflow as tf
import dlib
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.vgg16 import preprocess_input
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, BatchNormalization, Activation, Dropout

# Ruta al modelo de clasificación facial preentrenado
model_path = './face_classifier_model.h5'
# Ruta al modelo VGG-Face
vgg_face_path = './vgg_face_model.h5'

# Cargar el modelo VGG-Face preentrenado
vgg_face = load_model(vgg_face_path)

# Cargar el detector de rostros de Dlib basado en CNN
dnnFaceDetector = dlib.cnn_face_detection_model_v1("mmod_human_face_detector.dat")

# Cargar el modelo de clasificación sin compilar
classifier_model = tf.keras.models.load_model(model_path, compile=False)

# Compilar el modelo de clasificación con los parámetros adecuados
classifier_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                         optimizer='nadam',
                         metrics=['accuracy'])

# Diccionario que asigna los índices de las clases a nombres de personas
person_rep = {0: 'BradPitt',
 1: 'ScarlettJohansson',
 2: 'HughJackman',
 3: 'NataliePortman',
 4: 'TomHanks',
 5: 'KateWinslet',
 6: 'TomCruise',
 7: 'NicoleKidman',
 8: 'LeonardoDiCaprio',
 9: 'SandraBullock',
 10: 'WillSmith',
 11: 'AngelinaJolie',
 12: 'JenniferLawrence',
 13: 'JohnnyDepp',
 14: 'MeganFox',
 15: 'RobertDowneyJr',
 16: 'DenzelWashington'}

# Función para predecir la clase de una imagen
def predict_image(image_path):
    img = load_img(image_path, target_size=(224, 224))  # Cargar la imagen y redimensionarla
    img = img_to_array(img)  # Convertir la imagen a un array de numpy
    img = np.expand_dims(img, axis=0)  # Añadir una dimensión extra para el batch
    img = preprocess_input(img)  # Preprocesar la imagen para VGG-Face

    img_encode = vgg_face(img)  # Obtener los embeddings de la imagen
    predictions = classifier_model.predict(img_encode)  # Realizar la predicción
    return predictions

# Interfaz de usuario con Streamlit
st.title("Reconocimiento Facial")
st.write("Sube una imagen y el modelo intentará identificar el rostro.")

# Subida de archivo
uploaded_file = st.file_uploader("Elige una imagen...", type="jpg")

# Si se sube un archivo
if uploaded_file is not None:
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)  # Leer la imagen subida como un array de bytes
    img = cv2.imdecode(file_bytes, 1)  # Decodificar la imagen

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # Convertir la imagen a escala de grises
    rects = dnnFaceDetector(gray, 1)  # Detectar rostros en la imagen

    # Recorrer todos los rostros detectados
    for (i, rect) in enumerate(rects):
        left = rect.rect.left()  # Coordenada x1
        top = rect.rect.top()  # Coordenada y1
        right = rect.rect.right()  # Coordenada x2
        bottom = rect.rect.bottom()  # Coordenada y2
        width = right - left
        height = bottom - top
        img_crop = img[top:top + height, left:left + width]  # Recortar el rostro de la imagen
        cv2.imwrite('crop_img.jpg', img_crop)  # Guardar el rostro recortado temporalmente

        crop_img = load_img('crop_img.jpg', target_size=(224, 224))  # Cargar y redimensionar el rostro recortado
        crop_img = img_to_array(crop_img)  # Convertir el rostro a un array de numpy
        crop_img = np.expand_dims(crop_img, axis=0)  # Añadir una dimensión extra para el batch
        crop_img = preprocess_input(crop_img)  # Preprocesar la imagen para VGG-Face
        img_encode = vgg_face(crop_img)  # Obtener los embeddings del rostro

        embed = tf.keras.backend.eval(img_encode)  # Evaluar los embeddings
        person = classifier_model.predict(embed)  # Realizar la predicción
        confidence = np.max(person)  # Obtener la confianza de la predicción
        name = person_rep[np.argmax(person)]  # Obtener el nombre de la persona predicha
        os.remove('crop_img.jpg')  # Eliminar la imagen recortada temporalmente

        # Si la confianza es mayor a 0.9, dibujar el cuadro y el nombre en la imagen original
        if confidence > 0.9:
            cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 2)
            img = cv2.putText(img, name, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2, cv2.LINE_AA)
            img = cv2.putText(img, str(confidence), (left, bottom + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)

    # Mostrar la imagen procesada en la interfaz de Streamlit
    st.image(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), caption='Imagen procesada', use_column_width=True)