import depthai as dai
import cv2
import numpy as np

# Caminho para o modelo YOLOv7 convertido
MODEL_PATH = "seu_modelo_yolov7.blob"

# Criar pipeline
pipeline = dai.Pipeline()

# Criar um nó para a câmera
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setPreviewSize(640, 640)
cam_rgb.setInterleaved(False)
cam_rgb.setBoardSocket(dai.CameraBoardSocket.RGB)

# Criar um nó de Neural Network (NN) para executar o YOLOv7
nn = pipeline.create(dai.node.NeuralNetwork)
nn.setBlobPath(MODEL_PATH)
cam_rgb.preview.link(nn.input)

# Criar XLinkOut para visualizar os resultados
xout_video = pipeline.create(dai.node.XLinkOut)
xout_video.setStreamName("video")
cam_rgb.preview.link(xout_video.input)

xout_nn = pipeline.create(dai.node.XLinkOut)
xout_nn.setStreamName("detections")
nn.out.link(xout_nn.input)

# Rodar pipeline na OAK
with dai.Device(pipeline) as device:
    video_queue = device.getOutputQueue("video", maxSize=4, blocking=False)
    detection_queue = device.getOutputQueue("detections", maxSize=4, blocking=False)

    while True:
        frame = video_queue.get().getCvFrame()
        detections = detection_queue.tryGet()

        if detections is not None:
            for detection in detections.detections:
                x1, y1, x2, y2 = int(detection.xmin * 640), int(detection.ymin * 640), int(detection.xmax * 640), int(detection.ymax * 640)
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, f"Erva Daninha", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        cv2.imshow("YOLOv7 OAK-1 Lite W", frame)

        if cv2.waitKey(1) == ord('q'):
            break

cv2.destroyAllWindows()