agrobot_base/Python/OAK/imu.py

100 lines
3.8 KiB
Python

import depthai as dai
import time
import numpy as np
from ahrs.filters import Madgwick
from scipy.spatial.transform import Rotation as R
# ────────────────────────────────────────────────
# 🔧 Cria pipeline (RGB + IMU)
# ────────────────────────────────────────────────
pipeline = dai.Pipeline()
# 🔹 Câmera RGB
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setPreviewSize(300, 300)
camRgb.setInterleaved(False)
camRgb.setFps(30) # Taxa real da câmera
xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("rgb")
camRgb.preview.link(xoutRgb.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)
# ────────────────────────────────────────────────
# 🚀 Inicializa dispositivo
# ────────────────────────────────────────────────
with dai.Device(pipeline) as device:
print("Dispositivo conectado.")
rgbQueue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
imuQueue = device.getOutputQueue(name="imu", maxSize=50, blocking=False)
# 🔥 Filtro Madgwick
madgwick = Madgwick()
q = np.array([1.0, 0.0, 0.0, 0.0])
# ⏱️ FPS
frame_counter = 0
start_time = time.time()
last_log = time.time()
roll, pitch, yaw = 0.0, 0.0, 0.0
while True:
current_time = time.time()
# 🔸 IMU
imuData = imuQueue.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)
# 🔸 Leitura de frame RGB (apenas pra gerar FPS real)
frame = rgbQueue.tryGet()
if frame is not None:
frame_counter += 1
# 🔸 Faz log a cada 1 segundo
if (current_time - last_log) >= 1.0:
elapsed = current_time - start_time
fps = frame_counter / elapsed if elapsed > 0 else 0
# 🔥 Dados de inferência e memória são estimados (pode ser real no futuro)
tempo_inferencia_ms = 14.3
memoria_mib = 230
print("\n═══════════════════════════════════════════════")
print(f"📊 FPS da Pipeline : {fps:.2f} FPS")
print(f"⏱️ Tempo de inferência : {tempo_inferencia_ms:.2f} ms (simulado)")
print(f"🧠 Uso de memória VPU : {memoria_mib} MiB (estimado)")
print(f"🎯 Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°")
print("═══════════════════════════════════════════════")
# 🔄 Reseta contador
frame_counter = 0
start_time = time.time()
last_log = current_time