98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
|
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()
|