import cv2
import numpy as np

def detect_face_dnn(net, frame, conf_threshold=0.5):
    """
    Detecta rostros en una imagen utilizando un modelo DNN pre-entrenado
    y devuelve los cuadros delimitadores.
    """
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]
    
    # Crear un blob a partir de la imagen
    blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300), [104, 117, 123], False, False)
    
    # Establecer la entrada para la red neuronal
    net.setInput(blob)
    
    # Ejecutar la detección
    detections = net.forward()
    
    bboxes = []
    
    # Procesar las detecciones
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        
        # Filtrar por umbral de confianza
        if confidence > conf_threshold:
            x1 = int(detections[0, 0, i, 3] * frameWidth)
            y1 = int(detections[0, 0, i, 4] * frameHeight)
            x2 = int(detections[0, 0, i, 5] * frameWidth)
            y2 = int(detections[0, 0, i, 6] * frameHeight)
            
            # Asegurarse de que las coordenadas están dentro de los límites de la imagen
            x1 = max(0, min(x1, frameWidth - 1))
            y1 = max(0, min(y1, frameHeight - 1))
            x2 = max(0, min(x2, frameWidth - 1))
            y2 = max(0, min(y2, frameHeight - 1))
            
            # Agregar los cuadros delimitadores y la confianza
            bboxes.append([x1, y1, x2, y2, confidence])
    
    return np.array(bboxes)

def main():
    # Cargar el modelo de detección de rostros
    modelFile = "models/opencv_face_detector_uint8.pb"
    configFile = "models/opencv_face_detector.pbtxt"
    net = cv2.dnn.readNetFromTensorflow(modelFile, configFile)
    
    # Iniciar la cámara
    cap = cv2.VideoCapture(0)
    
    if not cap.isOpened():
        print("Error: No se pudo abrir la cámara.")
        return
    
    print("Cámara iniciada correctamente. Presione 'q' para salir.")
    
    while True:
        # Leer un frame de la cámara
        ret, frame = cap.read()
        
        if not ret:
            print("Error al capturar el frame.")
            break
        
        # Detectar rostros con umbrales de confianza diferentes
        for threshold in [0.5, 0.3, 0.1]:
            bboxes = detect_face_dnn(net, frame, threshold)
            
            # Crear copia del frame para dibujar en ella
            display_frame = frame.copy()
            
            # Dibujar los cuadros delimitadores
            for i, bbox in enumerate(bboxes):
                x1, y1, x2, y2, conf = bbox
                # Usar color verde
                color = (0, 255, 0)
                # Dibujar el rectángulo
                cv2.rectangle(display_frame, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
                # Mostrar la confianza
                label = f"Rostro: {conf:.2f}"
                cv2.putText(display_frame, label, (int(x1), int(y1) - 10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            
            # Mostrar el número de rostros detectados en la esquina
            cv2.putText(display_frame, f"Umbral: {threshold} - Rostros: {len(bboxes)}", 
                      (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            
            # Mostrar el resultado
            window_name = f"Detección de Rostros (Umbral: {threshold})"
            cv2.imshow(window_name, display_frame)
        
        # Salir si se presiona 'q'
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # Liberar recursos
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()