52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
|
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()
|