import cv2
import depthai as dai
import numpy as np
import argparse

# Configuração dos argumentos
parser = argparse.ArgumentParser(description="Rodar inferência na OAK-1 Lite W com YOLOv7")
parser.add_argument("--model", type=str, required=True, help="Caminho do modelo .blob")
args = parser.parse_args()

# Caminho do modelo
blob_path = args.model

# Criando pipeline
pipeline = dai.Pipeline()

# Criando um nó de câmera
cam = pipeline.create(dai.node.ColorCamera)
cam.setPreviewSize(640, 640)
cam.setInterleaved(False)
cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
cam.setFps(30)

# Criando um nó de detecção YOLO
nn = pipeline.create(dai.node.YoloDetectionNetwork)
nn.setBlobPath(blob_path)
nn.setConfidenceThreshold(0.5)  # Ajuste conforme necessário
nn.setNumClasses(10)  # Número de classes do modelo (atualize conforme necessário)
nn.setCoordinateSize(4)

# 🔹 Define âncoras e máscaras corretamente
nn.setAnchors([
    12, 16, 19, 36, 40, 28,  # Pequenos
    36, 75, 76, 55, 72, 146, # Médios
    142, 110, 192, 243, 459, 401 # Grandes
])
nn.setAnchorMasks({
    "side80": [1, 2, 3],  
    "side40": [3, 4, 5],  
    "side20": [5, 6, 7]   
})
nn.setIouThreshold(0.5)

# Conectando a câmera à rede neural
cam.preview.link(nn.input)

# Saída de vídeo para visualização
xout_cam = pipeline.create(dai.node.XLinkOut)
xout_cam.setStreamName("video")
cam.preview.link(xout_cam.input)

# Saída da inferência
xout_nn = pipeline.create(dai.node.XLinkOut)
xout_nn.setStreamName("detections")
nn.out.link(xout_nn.input)

# 🔹 Definição dos nomes das classes (ajuste conforme seu modelo)
labelMap = ['chenopodio', 'grama-azul', 'tiririca', 'erva-daninha', 'videira', 'mostarda-preta', 'milho', 'amaranta', 'farinha-seca', 'guanxuma']  # Substitua pelas classes corretas do seu modelo

# Iniciando dispositivo
with dai.Device(pipeline) as device:
    video_queue = device.getOutputQueue("video", maxSize=1, blocking=False)
    detections_queue = device.getOutputQueue("detections", maxSize=1, blocking=False)

    while True:
        # Captura a imagem
        frame = video_queue.get().getCvFrame()

        # Captura os resultados da inferência
        in_det = detections_queue.get()
        detections = in_det.detections

        print(f"Número de detecções: {len(detections)}")

        # Processando as detecções
        for detection in detections:
            x_min = int(detection.xmin * 640)
            y_min = int(detection.ymin * 640)
            x_max = int(detection.xmax * 640)
            y_max = int(detection.ymax * 640)
            conf = detection.confidence
            label_id = int(detection.label)  # ID da classe detectada

            if conf > 0.5:  # Ajuste o threshold de confiança conforme necessário
                label = labelMap[label_id] if label_id < len(labelMap) else f"Classe {label_id}"
                text = f"{label} ({conf:.2f})"

                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
                cv2.putText(frame, text, (x_min, y_min - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # Exibir a imagem
        cv2.imshow("Deteccao OAK", frame)
        if cv2.waitKey(1) == ord("q"):
            break

cv2.destroyAllWindows()