agrobot_base/Python/OAK/datasets/capture_worker.py

303 lines
12 KiB
Python

# capture_unificado.py com thread de status e IMU da OAK-D
import depthai as dai
import cv2
import os
import numpy as np
from datetime import datetime
import time
import threading
from ahrs.filters import Madgwick
from scipy.spatial.transform import Rotation as R
DIMENSAO_FINAL = (1280, 720)
QUALIDADE_JPEG = 95
TECLA_SALVAR = ord('s')
ALTURA_JANELA = 1080
LARGURA_JANELA = 1920
RECONNECT_INTERVAL = 5
AUTO_INTERVAL = 2.0
CAMERAS = {
"oak-d": None,
"oak-1": None
}
CAMERAS_IDS = {
"oak-d": "14442C10C143E2D600",
"oak-1": "14442C1011AD1ED000"
}
ULTIMA_TENTATIVA = {nome: 0 for nome in CAMERAS}
ULTIMA_CAPTURA = {nome: None for nome in CAMERAS}
STATUS_INFO = {nome: None for nome in CAMERAS}
IMU_INFO = {"oak-d": None}
AUTO_MODE = {nome: False for nome in CAMERAS}
ULTIMO_AUTO_SNAPSHOT = {nome: time.time() for nome in CAMERAS}
madgwick = Madgwick()
q_imu = np.array([1.0, 0.0, 0.0, 0.0])
PASTAS_SAIDA = {
"oak-d": os.path.join("oak-d", "dataset", "original", "images"),
"oak-1": os.path.join("oak-1", "dataset", "original", "images")
}
ISO_ATUAL = {
"oak-d": 400,
"oak-1": 400
}
EXPOSICAO_ATUAL = {
"oak-d": 6000,
"oak-1": 6000
}
ultimo_frame_salvo = {
"oak-d": np.zeros((360, 480, 3), dtype=np.uint8),
"oak-1": np.zeros((360, 480, 3), dtype=np.uint8)
}
camera_ativa = "oak-d"
for pasta in PASTAS_SAIDA.values():
os.makedirs(pasta, exist_ok=True)
def monitorar_status(nome):
while True:
time.sleep(2)
if CAMERAS[nome] is None:
STATUS_INFO[nome] = None
continue
try:
dev = CAMERAS[nome]["device"]
temp = dev.getChipTemperature().average
ddr = dev.getDdrMemoryUsage().used / 1024 / 1024
running = dev.isPipelineRunning()
speed = dev.getUsbSpeed().name
STATUS_INFO[nome] = {
"temp": temp,
"ddr": ddr,
"pipeline": running,
"usb": speed
}
except:
STATUS_INFO[nome] = None
def iniciar_dispositivo(nome):
try:
pipeline = dai.Pipeline()
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setPreviewSize(DIMENSAO_FINAL[0], DIMENSAO_FINAL[1])
camRgb.setInterleaved(False)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("rgb")
camRgb.preview.link(xoutRgb.input)
control_in = pipeline.create(dai.node.XLinkIn)
control_in.setStreamName("control")
control_in.out.link(camRgb.inputControl)
if nome == "oak-d":
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)
mx_id = CAMERAS_IDS[nome]
dev_info = dai.DeviceInfo(mx_id)
device = dai.Device(pipeline, dev_info)
rgbQueue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
controlQueue = device.getInputQueue("control")
cam = {
"device": device,
"queue": rgbQueue,
"control": controlQueue
}
if nome == "oak-d":
imuQueue = device.getOutputQueue(name="imu", maxSize=50, blocking=False)
cam["imu"] = imuQueue
CAMERAS[nome] = cam
return True
except Exception as e:
print(f"[{nome.upper()} ❌] Erro ao iniciar: {e}")
return False
def aplicar_controle_manual(nome):
if CAMERAS[nome] is None:
return
ctrl = dai.CameraControl()
ctrl.setManualExposure(EXPOSICAO_ATUAL[nome], ISO_ATUAL[nome])
ctrl.setManualFocus(130)
CAMERAS[nome]["control"].send(ctrl)
def monitorar_imu():
global q_imu
while True:
time.sleep(0.01)
cam = CAMERAS["oak-d"]
if cam is None or "imu" not in cam:
continue
imuQueue = cam["imu"]
imuData = imuQueue.tryGet()
if imuData is not None:
for packet in imuData.packets:
accel = packet.acceleroMeter
gyro = packet.gyroscope
ax, ay, az = accel.x, accel.y, accel.z
gx = np.deg2rad(gyro.x)
gy = np.deg2rad(gyro.y)
gz = np.deg2rad(gyro.z)
q_imu = madgwick.updateIMU(q=q_imu, gyr=np.array([gx, gy, gz]), acc=np.array([ax, ay, az]))
r = R.from_quat([q_imu[1], q_imu[2], q_imu[3], q_imu[0]])
roll, pitch, yaw = r.as_euler('xyz', degrees=True)
IMU_INFO["oak-d"] = {
"roll": roll,
"pitch": pitch,
"yaw": yaw
}
for nome in CAMERAS:
threading.Thread(target=monitorar_status, args=(nome,), daemon=True).start()
th_IMU = threading.Thread(target=monitorar_imu, daemon=True)
th_IMU.start()
while True:
canvas = np.zeros((ALTURA_JANELA, LARGURA_JANELA, 3), dtype=np.uint8)
agora = datetime.now()
timestamp_atual = agora.strftime("%d/%m/%Y %H:%M:%S")
agora_sec = time.time()
for nome in CAMERAS:
if CAMERAS[nome] is None and agora_sec - ULTIMA_TENTATIVA[nome] >= RECONNECT_INTERVAL:
print(f"[{nome.upper()} 🔄] Tentando reconectar...")
sucesso = iniciar_dispositivo(nome)
if sucesso:
aplicar_controle_manual(nome)
ULTIMA_TENTATIVA[nome] = agora_sec
for nome in CAMERAS:
if AUTO_MODE[nome] and (time.time() - ULTIMO_AUTO_SNAPSHOT[nome]) >= AUTO_INTERVAL:
if CAMERAS[nome] is not None:
try:
frame = CAMERAS[nome]["queue"].get().getCvFrame()
now = datetime.now()
timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
nome_base = f"img_{timestamp}.jpg"
caminho_final = os.path.join(PASTAS_SAIDA[nome], nome_base)
cv2.imwrite(caminho_final, frame, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
ultimo_frame_salvo[nome] = cv2.resize(frame, (480, 360))
ULTIMA_CAPTURA[nome] = now
print(f"[{nome.upper()} 🕒] Auto-snapshot salvo: {caminho_final}")
ULTIMO_AUTO_SNAPSHOT[nome] = time.time()
except Exception as e:
print(f"[{nome.upper()} ❌] Erro no auto-snapshot: {e}")
for idx, nome in enumerate(CAMERAS):
x_base = idx * 960
if CAMERAS[nome] is None:
cv2.putText(canvas, f"{nome.upper()} DESCONECTADA - Reconectando...", (x_base + 20, 50),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
continue
try:
in_rgb = CAMERAS[nome]["queue"].get()
frame = in_rgb.getCvFrame()
except Exception as e:
print(f"[{nome.upper()} ⚠️] Erro ao capturar frame: {e}")
CAMERAS[nome] = None
continue
frame_vivo = cv2.resize(frame, (960, 720))
captura_redimensionada = cv2.resize(ultimo_frame_salvo[nome], (480, 360))
canvas[0:720, x_base:x_base+960] = frame_vivo
canvas[720:1080, x_base:x_base+480] = captura_redimensionada
cor = (0,255,0) if camera_ativa == nome else (255,255,255)
cv2.putText(canvas, f"{nome.upper()} (Ativa: {'SIM' if camera_ativa==nome else 'NAO'}) | Modo: {'AUTO' if AUTO_MODE[nome] else 'MANUAL'}",
(x_base + 10, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, cor, 2)
cv2.putText(canvas, f"ISO: {ISO_ATUAL[nome]} | EXP: {EXPOSICAO_ATUAL[nome]}", (x_base + 10, 680), cv2.FONT_HERSHEY_SIMPLEX, 0.7, cor, 2)
cv2.putText(canvas, f"AO VIVO: {timestamp_atual}", (x_base + 10, 700), cv2.FONT_HERSHEY_SIMPLEX, 0.6, cor, 1)
if ULTIMA_CAPTURA[nome] is not None:
stamp = ULTIMA_CAPTURA[nome].strftime("%d/%m/%Y %H:%M:%S")
cv2.putText(canvas, f"ULTIMA CAPTURA: {stamp}", (x_base + 10, 740), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 1)
# Mostrar status cacheado se existir
info = STATUS_INFO.get(nome)
if info:
info_x = x_base + 500
y_base = 740
cv2.putText(canvas, f"Temp: {info['temp']:.2f} C", (info_x, y_base), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200,200,255), 1)
cv2.putText(canvas, f"DDR: {info['ddr']:.2f} MB", (info_x, y_base+30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200,200,255), 1)
cv2.putText(canvas, f"Pipeline: {'ON' if info['pipeline'] else 'OFF'}", (info_x, y_base+60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200,200,255), 1)
cv2.putText(canvas, f"USB: {info['usb']}", (info_x, y_base+90), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200,200,255), 1)
if nome == "oak-d" and IMU_INFO["oak-d"]:
imu = IMU_INFO["oak-d"]
cv2.putText(canvas, f"Roll: {imu['roll']:.2f}", (info_x, y_base+130), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180,255,180), 1)
cv2.putText(canvas, f"Pitch: {imu['pitch']:.2f}", (info_x, y_base+160), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180,255,180), 1)
cv2.putText(canvas, f"Yaw: {imu['yaw']:.2f}", (info_x, y_base+190), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180,255,180), 1)
if AUTO_MODE:
cv2.putText(canvas, f"AUTO SNAP: ON ({AUTO_INTERVAL:.1f}s)", (30, 1035), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,255), 2)
else:
cv2.putText(canvas, "AUTO SNAP: OFF", (30, 1035), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (100,100,100), 2)
cv2.putText(canvas, "Teclas: [1] OAK-D | [2] OAK-1 | s/+/-/m/n = camera ativa | q para sair",
(30, 1070), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (200, 200, 255), 2)
cv2.imshow("Captura Dual", canvas)
key = cv2.waitKey(1) & 0xFF
if key == ord('1'):
camera_ativa = "oak-d"
elif key == ord('2'):
camera_ativa = "oak-1"
elif key == TECLA_SALVAR:
if CAMERAS[camera_ativa] is not None:
try:
frame = CAMERAS[camera_ativa]["queue"].get().getCvFrame()
now = datetime.now()
timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
nome_base = f"img_{timestamp}.jpg"
caminho_final = os.path.join(PASTAS_SAIDA[camera_ativa], nome_base)
cv2.imwrite(caminho_final, frame, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
ultimo_frame_salvo[camera_ativa] = cv2.resize(frame, (480, 360))
ULTIMA_CAPTURA[camera_ativa] = now
print(f"[{camera_ativa.upper()} ✔] Imagem salva: {caminho_final}")
except Exception as e:
print(f"[{camera_ativa.upper()} ❌] Erro ao salvar imagem: {e}")
elif key == ord('+') or key == ord('='):
ISO_ATUAL[camera_ativa] = min(1600, ISO_ATUAL[camera_ativa] + 50)
aplicar_controle_manual(camera_ativa)
elif key == ord('-'):
ISO_ATUAL[camera_ativa] = max(100, ISO_ATUAL[camera_ativa] - 50)
aplicar_controle_manual(camera_ativa)
elif key == ord('m'):
EXPOSICAO_ATUAL[camera_ativa] = min(33000, EXPOSICAO_ATUAL[camera_ativa] + 1000)
aplicar_controle_manual(camera_ativa)
elif key == ord('n'):
EXPOSICAO_ATUAL[camera_ativa] = max(100, EXPOSICAO_ATUAL[camera_ativa] - 1000)
aplicar_controle_manual(camera_ativa)
elif key == ord('a'):
AUTO_MODE[camera_ativa] = not AUTO_MODE[camera_ativa]
print(f"[AUTO] {camera_ativa.upper()} modo automático {'ativado' if AUTO_MODE[camera_ativa] else 'desativado'}")
ULTIMO_AUTO_SNAPSHOT[camera_ativa] = time.time()
elif key == ord('q'):
print("[INFO] Encerrando...")
break
cv2.destroyAllWindows()