agrobot_base/Python/OAK/datasets/_0_capture.py

141 lines
5.1 KiB
Python

import depthai as dai
import cv2
import os
import numpy as np
from datetime import datetime
# === CONFIGURAÇÕES ===
MODELO = "oak-d"
PASTA_SAIDA = os.path.join(MODELO, "dataset")
DIMENSAO_FINAL = (1920, 1080)
QUALIDADE_JPEG = 95
TECLA_SALVAR = ord('s')
# Área de exibição
ALTURA_JANELA = 1080
LARGURA_JANELA = 1920
EXPOSICAO_ATUAL = [6000]
ISO_ATUAL = [400]
os.makedirs(PASTA_SAIDA, exist_ok=True)
# === PIPELINE ===
pipeline = dai.Pipeline()
cam_rgb = pipeline.createColorCamera()
cam_rgb.setPreviewSize(DIMENSAO_FINAL)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam_rgb.setInterleaved(False)
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
xout = pipeline.createXLinkOut()
xout.setStreamName("rgb")
cam_rgb.preview.link(xout.input)
control_in = pipeline.createXLinkIn()
control_in.setStreamName("control")
control_in.out.link(cam_rgb.inputControl)
# === EXECUÇÃO ===
with dai.Device(pipeline) as device:
print("[INFO] Pressione 's' para salvar | '+' e '-' para ISO | 'q' para sair.")
controlQueue = device.getInputQueue("control")
rgbQueue = device.getOutputQueue("rgb", maxSize=4, blocking=False)
def aplicar_controle_manual(exposicao, iso):
ctrl = dai.CameraControl()
ctrl.setManualExposure(exposicao, iso)
ctrl.setManualFocus(130)
controlQueue.send(ctrl)
print(f"[⚙️] Exposicao: {exposicao} µs | ISO: {iso}")
aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
ultimo_frame_salvo = np.zeros((DIMENSAO_FINAL[1], DIMENSAO_FINAL[0], 3), dtype=np.uint8)
# Fullscreen
cv2.namedWindow("Coletor de Dataset", cv2.WINDOW_NORMAL)
cv2.setWindowProperty("Coletor de Dataset", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
while True:
in_rgb = rgbQueue.get()
frame = in_rgb.getCvFrame()
# Redimensiona o frame ao vivo
frame_vivo = cv2.resize(frame, (1280, 960))
# Redimensiona a última captura
captura_redimensionada = cv2.resize(ultimo_frame_salvo, (512, 512))
# Cria o canvas vazio
canvas = np.zeros((ALTURA_JANELA, LARGURA_JANELA, 3), dtype=np.uint8)
# Posiciona o ao vivo à esquerda
canvas[60:1020, 30:1310] = frame_vivo
# Posiciona a última imagem salva à direita
canvas[60:572, 1400:1912] = captura_redimensionada
# Adiciona textos no lado direito
cv2.putText(canvas, f"ISO: {ISO_ATUAL[0]} | Exposicao: {EXPOSICAO_ATUAL[0]}", (1350, 620), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 255, 0), 2)
cv2.putText(canvas, "Pressione 's' para salvar", (1350, 680), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
cv2.putText(canvas, "'+' ou '-' para ajustar ISO", (1350, 720), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
cv2.putText(canvas, "'m' ou 'n' para ajustar a exposicao", (1350, 760), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
cv2.putText(canvas, "'q' para sair", (1350, 800), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 1)
cv2.imshow("Coletor de Dataset", canvas)
key = cv2.waitKey(1) & 0xFF
if key == TECLA_SALVAR:
now = datetime.now()
timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
caminho = os.path.join(PASTA_SAIDA, f"img_{timestamp}.jpg")
# Gera timestamp e nome-base
now = datetime.now()
timestamp = now.strftime("%Y%m%d_%H%M%S_%f")[:-3]
nome_base = f"img_{timestamp}.jpg"
# Original (sem resize)
os.makedirs(os.path.join("dataset", "original", "images"), exist_ok=True)
cv2.imwrite(os.path.join("dataset", "original", "images", nome_base), frame, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
# Demais resoluções
#resolucoes = {
# "512x512": (512, 512),
# "768x768": (768, 768),
# "384x384": (384, 384),
#}
#for nome_pasta, dim in resolucoes.items():
# path_pasta = os.path.join("dataset", nome_pasta)
# os.makedirs(path_pasta, exist_ok=True)
# imagem = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
# cv2.imwrite(os.path.join(path_pasta, nome_base), imagem, [cv2.IMWRITE_JPEG_QUALITY, QUALIDADE_JPEG])
print(f"[✔] Imagem salva: {caminho}")
ultimo_frame_salvo = frame.copy()
elif key == ord('+') or key == ord('='):
ISO_ATUAL[0] = min(1600, ISO_ATUAL[0] + 50)
aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
elif key == ord('-'):
ISO_ATUAL[0] = max(100, ISO_ATUAL[0] - 50)
aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
elif key == ord('m'):
EXPOSICAO_ATUAL[0] = min(10000, EXPOSICAO_ATUAL[0] + 500)
aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
elif key == ord('n'):
EXPOSICAO_ATUAL[0] = max(1000, EXPOSICAO_ATUAL[0] - 500)
aplicar_controle_manual(EXPOSICAO_ATUAL[0], ISO_ATUAL[0])
elif key == ord('q'):
print("[INFO] Encerrando...")
break
cv2.destroyAllWindows()