agrobot_base/Python/OAK/visao_geral.py

203 lines
6.9 KiB
Python

import depthai as dai
import cv2
import time
import numpy as np
from ahrs.filters import Madgwick
from scipy.spatial.transform import Rotation as R
# 🔥 Cria pipeline
pipeline = dai.Pipeline()
# 🔹 Câmera RGB
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setPreviewSize(640, 480)
camRgb.setInterleaved(False)
camRgb.setBoardSocket(dai.CameraBoardSocket.RGB)
camRgb.setFps(30)
xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("rgb")
camRgb.preview.link(xoutRgb.input)
# 🔹 Câmera Mono Left
monoLeft = pipeline.create(dai.node.MonoCamera)
monoLeft.setBoardSocket(dai.CameraBoardSocket.LEFT)
monoLeft.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
monoLeft.setFps(30)
xoutLeft = pipeline.create(dai.node.XLinkOut)
xoutLeft.setStreamName("left")
monoLeft.out.link(xoutLeft.input)
# 🔹 Câmera Mono Right
monoRight = pipeline.create(dai.node.MonoCamera)
monoRight.setBoardSocket(dai.CameraBoardSocket.RIGHT)
monoRight.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
monoRight.setFps(30)
xoutRight = pipeline.create(dai.node.XLinkOut)
xoutRight.setStreamName("right")
monoRight.out.link(xoutRight.input)
# 🔹 Profundidade
stereo = pipeline.create(dai.node.StereoDepth)
monoLeft.out.link(stereo.left)
monoRight.out.link(stereo.right)
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_ACCURACY)
xoutDepth = pipeline.create(dai.node.XLinkOut)
xoutDepth.setStreamName("depth")
stereo.depth.link(xoutDepth.input)
# 🔹 IMU
imu = pipeline.create(dai.node.IMU)
imu.enableIMUSensor(dai.IMUSensor.ACCELEROMETER_RAW, 500)
imu.enableIMUSensor(dai.IMUSensor.GYROSCOPE_RAW, 500)
imu.setBatchReportThreshold(1)
imu.setMaxBatchReports(20)
xoutImu = pipeline.create(dai.node.XLinkOut)
xoutImu.setStreamName("imu")
imu.out.link(xoutImu.input)
segmentation = pipeline.create(dai.node.NeuralNetwork)
segmentation.setBlobPath("model.blob")
camRgb.preview.link(segmentation.input)
xoutNN = pipeline.create(dai.node.XLinkOut)
xoutNN.setStreamName("nn")
segmentation.out.link(xoutNN.input)
# 🚀 Executa
with dai.Device(pipeline) as device:
print("Dispositivo conectado.")
qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
qLeft = device.getOutputQueue(name="left", maxSize=4, blocking=False)
qRight = device.getOutputQueue(name="right", maxSize=4, blocking=False)
qDepth = device.getOutputQueue(name="depth", maxSize=4, blocking=False)
qImu = device.getOutputQueue(name="imu", maxSize=50, blocking=False)
# 🧠 Filtro Madgwick
madgwick = Madgwick()
q = np.array([1.0, 0.0, 0.0, 0.0])
roll, pitch, yaw = 0.0, 0.0, 0.0
frame_counter = 0
start_time = time.time()
last_log = time.time()
while True:
current_time = time.time()
# 🔸 IMU
imuData = qImu.tryGet()
if imuData is not None:
for packet in imuData.packets:
accel = packet.acceleroMeter
gyro = packet.gyroscope
ax = accel.x
ay = accel.y
az = accel.z
gx = np.deg2rad(gyro.x)
gy = np.deg2rad(gyro.y)
gz = np.deg2rad(gyro.z)
q = madgwick.updateIMU(q=q, gyr=np.array([gx, gy, gz]), acc=np.array([ax, ay, az]))
r = R.from_quat([q[1], q[2], q[3], q[0]])
roll, pitch, yaw = r.as_euler('xyz', degrees=True)
# 🔸 Frames
inRgb = qRgb.tryGet()
inLeft = qLeft.tryGet()
inRight = qRight.tryGet()
inDepth = qDepth.tryGet()
nnQueue = device.getOutputQueue(name="nn", maxSize=4, blocking=False)
in_nn = nnQueue.tryGet()
if in_nn is not None:
data = in_nn.getFirstLayerFp16()
# 🔥 Ajuste pro tamanho do modelo que você tá usando
H = 512
W = 512
num_classes = 3 # Se seu modelo tiver 3 classes, ajusta conforme seu modelo
data = np.array(data).reshape((num_classes, H, W))
mask = np.argmax(data, axis=0).astype(np.uint8)
# 🔥 Cores por classe
color_map = {
0: (0, 0, 0), # Fundo (preto)
1: (0, 255, 0), # Planta (verde)
2: (255, 0, 0) # Céu (azul)
}
mask_color = np.zeros((H, W, 3), dtype=np.uint8)
for class_id, color in color_map.items():
mask_color[mask == class_id] = color
# 🔥 Redimensiona a máscara pra imagem RGB
mask_color = cv2.resize(mask_color, (imgRgb.shape[1], imgRgb.shape[0]))
# 🔥 Overlay na imagem RGB
overlay = cv2.addWeighted(imgRgb, 0.6, mask_color, 0.4, 0)
cv2.imshow("Segmentacao", overlay)
if inRgb: frame_counter += 1
# 🔥 Monta janela
imgRgb = inRgb.getCvFrame() if inRgb else np.zeros((480, 640, 3), dtype=np.uint8)
imgLeft = inLeft.getCvFrame() if inLeft else np.zeros((400, 640), dtype=np.uint8)
imgRight = inRight.getCvFrame() if inRight else np.zeros((400, 640), dtype=np.uint8)
imgDepth = inDepth.getFrame() if inDepth else np.zeros((400, 640), dtype=np.uint16)
if imgDepth.max() > 0:
imgDepthNorm = cv2.normalize(imgDepth, None, 0, 255, cv2.NORM_MINMAX)
imgDepthNorm = imgDepthNorm.astype(np.uint8)
imgDepthColor = cv2.applyColorMap(imgDepthNorm, cv2.COLORMAP_JET)
else:
imgDepthColor = np.zeros((400, 640, 3), dtype=np.uint8)
# 🔥 Mostra imagens
cv2.imshow("RGB", imgRgb)
cv2.imshow("Mono Left", imgLeft)
cv2.imshow("Mono Right", imgRight)
cv2.imshow("Depth", imgDepthColor)
# 🔥 Logs a cada segundo
if (current_time - last_log) >= 1.0:
elapsed = current_time - start_time
fps = frame_counter / elapsed if elapsed > 0 else 0
try:
mem = device.getDdrMemoryUsage()
temp = device.getChipTemperature()
except:
mem = None
temp = None
print("\n═══════════════════════════════════════════════")
print(f"📊 FPS da Pipeline : {fps:.2f} FPS")
print(f"🌡️ Temp CSS:{temp.css:.2f}° MSS:{temp.mss:.2f}° DSS:{temp.dss:.2f}° UPA:{temp.upa:.2f}°")
print(f"🧠 Memória DDR : {mem.used} MiB usada / {mem.total} MiB total")
print(f"🎯 Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°")
print("═══════════════════════════════════════════════")
frame_counter = 0
start_time = time.time()
last_log = current_time
# 🔴 Encerra com Q
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()