import cv2

# Carregar o modelo pré-treinado e os arquivos de configuração
model_path = 'C:/Zendion Inc/agrobot_base/Python/models/coco/faster_rcnn_inception_v2_coco_2018_01_28/frozen_inference_graph.pb'
config_path = 'C:/Zendion Inc/agrobot_base/Python/models/coco/faster_rcnn_inception_v2_coco_2018_01_28.pbtxt'

# Carregar o modelo
net = cv2.dnn_DetectionModel(model_path, config_path)
net.setInputSize(320, 320)  # Definir o tamanho da entrada do modelo

# Definir as classes que o modelo pode detectar
classNames = []
with open('models/ssd_mobilenet_v3_large_coco/coco.names', 'rt') as f:
    classNames = f.read().rstrip('\n').split('\n')

# Configurações adicionais
conf_threshold = 0.6  # Limiar de confiança para detecção

# 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

    # Detectar objetos no frame
    classIds, confs, bbox = net.detect(frame, confThreshold=conf_threshold)

    if len(classIds) > 0:
        for classId, confidence, box in zip(classIds.flatten(), confs.flatten(), bbox):
            cv2.rectangle(frame, box, color=(0, 255, 0), thickness=2)
            cv2.putText(frame, classNames[classId - 1], (box[0] + 10, box[1] + 30),
                        cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)

    cv2.imshow('Detecção de Objetos', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()