import cv2
import numpy as np

# Carregar o modelo YOLO e os arquivos de configuração
model_config = 'C:/Zendion Inc/agrobot_base/Python/models/yolo/yolov3.cfg'
model_weights = 'C:/Zendion Inc/agrobot_base/Python/models/yolo/yolov3.weights'

# Carregar as classes que o modelo pode detectar
with open('C:/Zendion Inc/agrobot_base/Python/models/yolo/coco.names', 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

# Carregar o modelo YOLO
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)  # Use DNN_TARGET_OPENCL para GPU

# Configurações adicionais
conf_threshold = 0.5
nms_threshold = 0.4

# Iniciar a captura de vídeo (altere o índice conforme necessário)
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # Obter as camadas de saída do modelo YOLO
    output_layers_names = net.getUnconnectedOutLayersNames()

    # Verificar se há saída válida
    if not output_layers_names:
        print("Nenhuma camada de saída encontrada. Verifique se o modelo foi carregado corretamente.")
        break

    # Fazer a detecção de objetos no frame
    blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
    net.setInput(blob)
    outs = net.forward(output_layers_names)

    # Mostrar as detecções no frame
    class_ids = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]

            if confidence > conf_threshold:
                # Obter informações do objeto detectado
                center_x = int(detection[0] * frame.shape[1])
                center_y = int(detection[1] * frame.shape[0])
                w = int(detection[2] * frame.shape[1])
                h = int(detection[3] * frame.shape[0])

                # Coordenadas do retângulo
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)

                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)

    # Aplicar Non-Max Suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # Mostrar as detecções no frame após o NMS
    for i in indices:
        box = boxes[i]
        x, y, w, h = box[0], box[1], box[2], box[3]
        cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
        cv2.putText(frame, classes[class_ids[i]], (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    cv2.imshow('Deteccao de Objetos YOLO', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()