112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
|
import os
|
||
|
|
import depthai as dai
|
||
|
|
import numpy as np
|
||
|
|
import cv2
|
||
|
|
import time
|
||
|
|
from utils import converter_mask_ids_para_rgb, carregar_labelmap_completo
|
||
|
|
|
||
|
|
# ⚙️ Configurações
|
||
|
|
MODELO = "oak-1"
|
||
|
|
MODEL_NAME = "ervas_full"
|
||
|
|
RESOLUCAO = (384, 384)
|
||
|
|
ROI_INICIO = 0.0
|
||
|
|
ROI_TAMANHO = 1.0
|
||
|
|
blob_path = os.path.join(MODELO, "backup", "fast_scnn", MODEL_NAME, MODEL_NAME + "_best_openvino_2022.1_6shave.blob")
|
||
|
|
labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")
|
||
|
|
model_name = MODEL_NAME + "_best"
|
||
|
|
|
||
|
|
# Carregar mapa de cores
|
||
|
|
_, colormap_rgb, classes, ignore_rgb = carregar_labelmap_completo(labelmap_path)
|
||
|
|
NUM_CLASSES = len(classes)
|
||
|
|
IGNORE_ID = ignore_rgb[0]
|
||
|
|
|
||
|
|
# Criar pipeline
|
||
|
|
pipeline = dai.Pipeline()
|
||
|
|
|
||
|
|
# Câmera
|
||
|
|
cam = pipeline.createColorCamera()
|
||
|
|
cam.setBoardSocket(dai.CameraBoardSocket.CAM_A)
|
||
|
|
cam.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
|
||
|
|
cam.setInterleaved(False)
|
||
|
|
cam.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)
|
||
|
|
cam.setFps(30)
|
||
|
|
|
||
|
|
# ImageManip (redimensiona para a entrada da rede)
|
||
|
|
manip = pipeline.createImageManip()
|
||
|
|
|
||
|
|
y1 = 1.0 - (ROI_INICIO + ROI_TAMANHO)
|
||
|
|
y2 = 1.0 - ROI_INICIO
|
||
|
|
|
||
|
|
manip.initialConfig.setCropRect(0.0, y1, 1.0, y2)
|
||
|
|
manip.initialConfig.setResize(RESOLUCAO[1], RESOLUCAO[0])
|
||
|
|
|
||
|
|
manip.initialConfig.setFrameType(dai.RawImgFrame.Type.RGB888p)
|
||
|
|
cam.preview.link(manip.inputImage)
|
||
|
|
|
||
|
|
# Neural network
|
||
|
|
nn = pipeline.createNeuralNetwork()
|
||
|
|
nn.setBlobPath(blob_path)
|
||
|
|
manip.out.link(nn.input)
|
||
|
|
|
||
|
|
# Saída RGB para overlay (sem redimensionar)
|
||
|
|
xout_rgb = pipeline.createXLinkOut()
|
||
|
|
xout_rgb.setStreamName("rgb")
|
||
|
|
cam.preview.link(xout_rgb.input)
|
||
|
|
|
||
|
|
# Saída NN
|
||
|
|
xout_nn = pipeline.createXLinkOut()
|
||
|
|
xout_nn.setStreamName("nn")
|
||
|
|
nn.out.link(xout_nn.input)
|
||
|
|
|
||
|
|
# Rodar pipeline
|
||
|
|
with dai.Device(pipeline) as device:
|
||
|
|
rgb_queue = device.getOutputQueue("rgb", maxSize=1, blocking=False)
|
||
|
|
nn_queue = device.getOutputQueue("nn", maxSize=1, blocking=False)
|
||
|
|
|
||
|
|
print("Rodando inferência na OAK... Pressione 'q' para sair.")
|
||
|
|
prev_time = time.time()
|
||
|
|
|
||
|
|
while True:
|
||
|
|
in_rgb = rgb_queue.get()
|
||
|
|
in_nn = nn_queue.get()
|
||
|
|
|
||
|
|
# RGB frame da câmera
|
||
|
|
frame = in_rgb.getCvFrame()
|
||
|
|
|
||
|
|
# Inferência - saída é um vetor flat [num_classes * H * W]
|
||
|
|
out = in_nn.getFirstLayerFp16()
|
||
|
|
out_np = np.array(out, dtype=np.float32).reshape((NUM_CLASSES, RESOLUCAO[1], RESOLUCAO[0]))
|
||
|
|
|
||
|
|
# Pega o índice da classe com maior probabilidade por pixel
|
||
|
|
pred_ids = np.argmax(out_np, axis=0).astype(np.uint8)
|
||
|
|
|
||
|
|
# Converter para RGB bonitão
|
||
|
|
pred_rgb = converter_mask_ids_para_rgb(pred_ids, colormap_rgb, IGNORE_ID)
|
||
|
|
roi_h = int((y2 - y1) * frame.shape[0])
|
||
|
|
roi_w = frame.shape[1]
|
||
|
|
pred_rgb_resized = cv2.resize(pred_rgb, (roi_w, roi_h), interpolation=cv2.INTER_NEAREST)
|
||
|
|
y_start = int(y1 * frame.shape[0])
|
||
|
|
y_end = y_start + roi_h
|
||
|
|
y_start = max(0, min(frame.shape[0], y_start))
|
||
|
|
y_end = max(0, min(frame.shape[0], y_end))
|
||
|
|
overlay = frame.copy()
|
||
|
|
overlay[y_start:y_end, 0:roi_w] = cv2.addWeighted(
|
||
|
|
frame[y_start:y_end, 0:roi_w], 0.4, pred_rgb_resized, 0.6, 0
|
||
|
|
)
|
||
|
|
|
||
|
|
# FPS
|
||
|
|
now = time.time()
|
||
|
|
fps = 1.0 / (now - prev_time)
|
||
|
|
prev_time = now
|
||
|
|
cv2.putText(overlay, f"FPS: {fps:.1f}", (10, 30),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
||
|
|
|
||
|
|
# Redimensiona para tela cheia (por exemplo 1280x720 ou tela do usuário)
|
||
|
|
overlay_display = cv2.resize(overlay, (1280, 720))
|
||
|
|
cv2.imshow("Segmentação - OAK (on-board)", cv2.cvtColor(overlay_display, cv2.COLOR_RGB2BGR))
|
||
|
|
|
||
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||
|
|
break
|
||
|
|
|
||
|
|
cv2.destroyAllWindows()
|