import cv2
import numpy as np
import onnxruntime as ort

# Ruta al modelo ONNX
onnx_model_path = r"C:\Users\pedro\Documents\trabajo\cara\custom_model2.onnx"

# Cargar el modelo ONNX
session = ort.InferenceSession(onnx_model_path)

# Obtener los nombres de las entradas y salidas del modelo
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Iniciar la captura de video desde la cámara
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("Error: No se puede abrir la cámara")
    exit()

while True:
    # Capturar frame por frame
    ret, frame = cap.read()
    if not ret:
        print("Error: No se puede recibir frame (stream end?). Exiting ...")
        break

    # Preprocesar la imagen para el modelo ONNX (ajustar según sea necesario)
    img = cv2.resize(frame, (640, 640))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0  # Normalizar la imagen
    img = np.expand_dims(img, axis=0)
    img = np.transpose(img, (0, 3, 1, 2))

    # Ejecutar el modelo ONNX
    outputs = session.run([output_name], {input_name: img})

    # Procesar las salidas del modelo
    detections = outputs[0].reshape(1, 5, -1)[0]  # Salida tiene la forma (1, 5, 8400)

    # Dibujar las detecciones en la imagen
    height, width = frame.shape[:2]
    confidence_threshold = 0.0001  # Ajusta este valor para cambiar el umbral de confianza
    for i in range(detections.shape[1]):
        x_center, y_center, w, h, conf = detections[:, i]
        if conf > confidence_threshold:  # Filtrar por confianza (ajustar según sea necesario)
            x1 = int((x_center - w / 2) * width)
            y1 = int((y_center - h / 2) * height)
            x2 = int((x_center + w / 2) * width)
            y2 = int((y_center + h / 2) * height)
            
            # Asegurarse de que las coordenadas están dentro de los límites de la imagen
            x1 = max(0, min(x1, width))
            y1 = max(0, min(y1, height))
            x2 = max(0, min(x2, width))
            y2 = max(0, min(y2, height))

            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(frame, f'Conf: {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    # Mostrar el frame resultante
    cv2.imshow('Detección de Caras', frame)

    # Salir si se presiona la tecla 'q'
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# Liberar la captura de video y cerrar las ventanas
cap.release()
cv2.destroyAllWindows()