203 lines
6.9 KiB
Python
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()
|