328 lines
12 KiB
Python
328 lines
12 KiB
Python
import os
|
|
import time
|
|
import json
|
|
import argparse
|
|
from datetime import datetime
|
|
|
|
import numpy as np
|
|
import cv2
|
|
|
|
from raw_segformer_service import make_bgr_preview_from_raw
|
|
from gal5000.gal_service import Gal5000Camera # ajuste o nome do módulo se estiver diferente
|
|
|
|
|
|
# =========================
|
|
# Helpers gerais
|
|
# =========================
|
|
|
|
def clamp(v, lo, hi):
|
|
return lo if v < lo else hi if v > hi else v
|
|
|
|
|
|
def ts_name() -> str:
|
|
"""Timestamp legível e único para nome de arquivo."""
|
|
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
|
|
|
|
|
def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
|
|
"""
|
|
Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
|
|
Pensado pra deixar o preview bonitinho sem estourar tudo.
|
|
"""
|
|
x = np.asarray(x)
|
|
if x.dtype != np.float32 and x.dtype != np.float64:
|
|
x = x.astype(np.float32)
|
|
|
|
# Se o canal já está em 0..1, escala pra 0..255 antes de cortar
|
|
if x.max() <= 1.5:
|
|
x = x * 255.0
|
|
|
|
lo = np.percentile(x, p_lo)
|
|
hi = np.percentile(x, p_hi)
|
|
|
|
if hi <= lo + 1e-3:
|
|
y = x
|
|
else:
|
|
y = (x - lo) * (255.0 / (hi - lo))
|
|
|
|
return np.clip(y, 0, 255).astype(np.uint8)
|
|
|
|
|
|
def overlay_hud(
|
|
img_bgr: np.ndarray,
|
|
lines: list[str],
|
|
base_h: int = 720,
|
|
base_font_scale: float = 0.75,
|
|
base_line_step: int = 28,
|
|
):
|
|
"""
|
|
Escreve textos empilhados no canto superior esquerdo,
|
|
ajustando o tamanho do texto de acordo com a altura da imagem.
|
|
|
|
base_h: altura de referência (ex: 720 ou a RAW_H original).
|
|
"""
|
|
h, w = img_bgr.shape[:2]
|
|
|
|
# Fator de escala com base na altura atual
|
|
scale = h / float(base_h)
|
|
|
|
# Evita ficar microscópico em resoluções muito baixas
|
|
scale = max(scale, 0.4)
|
|
|
|
font_scale = base_font_scale * scale
|
|
line_step = int(base_line_step * scale)
|
|
|
|
# Espessuras proporcionais
|
|
thick_outline = max(1, int(3 * scale))
|
|
thick_text = max(1, int(2 * scale))
|
|
|
|
# Margem superior / esquerda também escaladas
|
|
y = int(24 * scale)
|
|
x = int(12 * scale)
|
|
|
|
for s in lines:
|
|
# contorno preto
|
|
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
|
|
# texto branco
|
|
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
|
|
y += line_step
|
|
|
|
|
|
def save_sample_raw4(
|
|
base_dir: str,
|
|
raw4: np.ndarray,
|
|
preview_bgr: np.ndarray,
|
|
meta: dict,
|
|
):
|
|
"""
|
|
Salva:
|
|
- RAW4 como .raw float32 (4,H,W)
|
|
- preview RGB como .png
|
|
- metadados como .json
|
|
dentro de base_dir.
|
|
"""
|
|
os.makedirs(base_dir, exist_ok=True)
|
|
name = ts_name()
|
|
|
|
raw_path = os.path.join(base_dir, f"{name}.raw")
|
|
png_path = os.path.join(base_dir, f"{name}.png")
|
|
json_path = os.path.join(base_dir, f"{name}.json")
|
|
|
|
# RAW4
|
|
#np.save(raw_path, raw4.astype(np.float32))
|
|
raw4.astype(np.float32).tofile(raw_path)
|
|
|
|
# Preview
|
|
cv2.imwrite(png_path, preview_bgr)
|
|
|
|
# Metadados
|
|
with open(json_path, "w", encoding="utf-8") as f:
|
|
json.dump(meta, f, ensure_ascii=False, indent=2)
|
|
|
|
return raw_path, png_path, json_path
|
|
|
|
|
|
# =========================
|
|
# MAIN
|
|
# =========================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
|
|
parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
|
|
parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
|
|
parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
|
|
parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
|
|
parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
|
|
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
|
|
parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
|
|
parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# ===== config =====
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
MODELO = config["camera"]
|
|
RAW_W = config["raw_size"][0]
|
|
RAW_H = config["raw_size"][1]
|
|
|
|
# Define diretório de sessão:
|
|
# dataset/cana_<estado>/<horario>/<YYYYMMDD>/
|
|
session_dir = os.path.join(MODELO, args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
|
|
os.makedirs(session_dir, exist_ok=True)
|
|
|
|
print("============================================")
|
|
print("Coleta de dataset RAW4 - SegFormer B0")
|
|
print(f"Cana : {args.cana}")
|
|
print(f"Horário : {args.horario}")
|
|
print(f"Saída : {session_dir}")
|
|
print("============================================")
|
|
|
|
window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
|
|
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
|
|
|
auto_save = False
|
|
last_auto_t = 0.0
|
|
upscale = args.upscale
|
|
|
|
# Estatísticas simples de FPS
|
|
t_fps = time.time()
|
|
frames = 0
|
|
fps = 0.0
|
|
|
|
last_msg = ""
|
|
last_msg_t = 0.0
|
|
|
|
try:
|
|
cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
|
|
with cam:
|
|
print("[CAM] Status inicial:", cam.get_status())
|
|
|
|
# opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
|
|
cam.configure_fps(20)
|
|
cam.start_streaming()
|
|
|
|
while True:
|
|
t0 = time.time()
|
|
raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
|
|
t1 = time.time()
|
|
|
|
ae_dbg = dbg.get("ae", {}) or {}
|
|
exp_raw = dbg.get("exp_raw", None)
|
|
gain_a = dbg.get("gain_a", None)
|
|
gain_d = dbg.get("gain_d", None)
|
|
|
|
apply_ir_comp = True
|
|
ir_k_r = 0.4
|
|
ir_k_g = 0.1
|
|
ir_k_b = 0.5
|
|
bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
|
|
|
# FPS
|
|
frames += 1
|
|
dt_fps = time.time() - t_fps
|
|
if dt_fps >= 1.0:
|
|
fps = frames / dt_fps
|
|
frames = 0
|
|
t_fps = time.time()
|
|
|
|
# AE info
|
|
p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
|
|
sat_disp = ae_dbg.get("sat", 0.0)
|
|
hold = ae_dbg.get("hold", False)
|
|
|
|
ae_on = cam.is_auto_exposure_enabled()
|
|
|
|
# HUD principal
|
|
lines = [
|
|
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
|
f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
|
|
f"exp_raw={exp_raw} gain_a={gain_a} gain_d={gain_d} | FPS={fps:.1f}",
|
|
f"AEdbg: p95={p95_disp:.1f} sat={sat_disp:.3f} hold={hold}",
|
|
"Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
|
|
]
|
|
overlay_hud(bgr, lines, base_h=RAW_H)
|
|
|
|
# Mensagem rápida (ex: arquivo salvo)
|
|
if last_msg and (time.time() - last_msg_t) < 2.0:
|
|
cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
|
|
|
cv2.imshow(window_name, bgr)
|
|
|
|
# Auto-save
|
|
now = time.time()
|
|
if auto_save and (now - last_auto_t) >= args.interval:
|
|
meta = {
|
|
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
|
"cana": args.cana,
|
|
"horario": args.horario,
|
|
"raw4_shape": list(raw4_base.shape),
|
|
"out_h": RAW_H,
|
|
"out_w": RAW_W,
|
|
"ae_enabled": bool(ae_on),
|
|
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
|
"gain_a": int(gain_a) if gain_a is not None else None,
|
|
"gain_d": int(gain_d) if gain_d is not None else None,
|
|
"apply_ir_comp": apply_ir_comp,
|
|
"ir_k_r": ir_k_r,
|
|
"ir_k_g": ir_k_g,
|
|
"ir_k_b": ir_k_b,
|
|
"ae_dbg": {
|
|
k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
|
|
for k, v in ae_dbg.items()
|
|
},
|
|
"note": "autosave",
|
|
}
|
|
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False, apply_ir_comp=apply_ir_comp, ir_k_r=ir_k_r, ir_k_g=ir_k_g, ir_k_b=ir_k_b)
|
|
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
|
last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
|
|
last_msg_t = now
|
|
last_auto_t = now
|
|
|
|
# Teclado
|
|
k = cv2.waitKey(1) & 0xFF
|
|
if k in (ord("q"), ord("Q"), 27): # Q ou ESC
|
|
break
|
|
|
|
elif k in (ord("a"), ord("A")):
|
|
auto_save = not auto_save
|
|
last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
|
|
last_msg_t = time.time()
|
|
|
|
elif k in (ord("e"), ord("E")):
|
|
cam.enable_auto_exposure(not ae_on)
|
|
last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
|
|
last_msg_t = time.time()
|
|
|
|
elif k in (ord("m"), ord("M")):
|
|
upscale = 0 if upscale else args.upscale
|
|
last_msg = f"Preview UPSCALE -> {upscale}"
|
|
last_msg_t = time.time()
|
|
|
|
elif k in (ord("c"), ord("C"), 32): # C ou SPACE
|
|
meta = {
|
|
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
|
"cana": args.cana,
|
|
"horario": args.horario,
|
|
"raw4_shape": list(raw4_base.shape),
|
|
"out_h": RAW_H,
|
|
"out_w": RAW_W,
|
|
"ae_enabled": bool(ae_on),
|
|
"exp_raw": int(exp_raw) if exp_raw is not None else None,
|
|
"gain_a": int(gain_a) if gain_a is not None else None,
|
|
"gain_d": int(gain_d) if gain_d is not None else None,
|
|
"apply_ir_comp": apply_ir_comp,
|
|
"ir_k_r": ir_k_r,
|
|
"ir_k_g": ir_k_g,
|
|
"ir_k_b": ir_k_b,
|
|
"ae_dbg": {
|
|
k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
|
|
for k2, v2 in ae_dbg.items()
|
|
},
|
|
"note": "manual",
|
|
}
|
|
rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False)
|
|
raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
|
|
last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
|
|
last_msg_t = time.time()
|
|
|
|
# você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)
|
|
|
|
# Só pra não ficar rodando a 1000 FPS na UI
|
|
# mas sem travar muito a captura
|
|
dt_loop = time.time() - t0
|
|
if dt_loop < 0.001:
|
|
time.sleep(0.001)
|
|
|
|
finally:
|
|
cv2.destroyAllWindows()
|
|
print("Fim da captura.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|