497 lines
22 KiB
Python
497 lines
22 KiB
Python
import os
|
|
import time
|
|
import json
|
|
import argparse
|
|
from datetime import datetime
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from core.oak_fcc3_client import OakFcc3Client as MultiSpectralClient
|
|
|
|
with open("config.json", "r", encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
|
|
RAW_SIZE = config.get("raw_size") # [W, H]
|
|
MODULE_PARAMS = config.get("module_params_json")
|
|
|
|
|
|
# =========================
|
|
# Helpers gerais
|
|
# =========================
|
|
|
|
def ts_name() -> str:
|
|
return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
|
|
|
|
|
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,
|
|
):
|
|
h, w = img_bgr.shape[:2]
|
|
|
|
scale = h / float(base_h)
|
|
scale = max(scale, 0.4)
|
|
|
|
font_scale = base_font_scale * scale
|
|
line_step = int(base_line_step * scale)
|
|
|
|
thick_outline = max(1, int(3 * scale))
|
|
thick_text = max(1, int(2 * scale))
|
|
|
|
y = int(24 * scale)
|
|
x = int(12 * scale)
|
|
|
|
for s in lines:
|
|
cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
|
|
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(
|
|
base_dir: str,
|
|
frame_type: str,
|
|
preview_bgr: np.ndarray,
|
|
meta: dict,
|
|
raw_payload: np.ndarray | None = None,
|
|
packed_raw: np.ndarray | None = None,
|
|
packed_raw_by_camera: dict | None = None,
|
|
):
|
|
os.makedirs(base_dir, exist_ok=True)
|
|
name = ts_name()
|
|
|
|
png_path = os.path.join(base_dir, f"{name}.png")
|
|
json_path = os.path.join(base_dir, f"{name}.json")
|
|
|
|
if frame_type in ("RGB", "MULTISPEC"):
|
|
if raw_payload is None:
|
|
raise ValueError(f"raw_payload não pode ser None quando frame_type='{frame_type}'")
|
|
|
|
payload_path = os.path.join(base_dir, f"{name}.raw")
|
|
raw_payload.astype(np.float32).tofile(payload_path)
|
|
|
|
meta["saved_payload_type"] = frame_type.lower()
|
|
meta["saved_payload_path"] = os.path.basename(payload_path)
|
|
meta["saved_payload_dtype"] = "float32"
|
|
meta["saved_payload_shape"] = list(raw_payload.shape)
|
|
|
|
elif frame_type == "RAW_BRUTO":
|
|
if packed_raw_by_camera is not None:
|
|
payload_files = {}
|
|
payload_shapes = {}
|
|
payload_dtypes = {}
|
|
|
|
for cam_id, arr in packed_raw_by_camera.items():
|
|
path = os.path.join(base_dir, f"{name}_{cam_id}.bin")
|
|
arr.tofile(path)
|
|
payload_files[cam_id] = os.path.basename(path)
|
|
payload_shapes[cam_id] = list(arr.shape)
|
|
payload_dtypes[cam_id] = str(arr.dtype)
|
|
|
|
meta["saved_payload_type"] = "raw_native_multi"
|
|
meta["saved_payload_paths"] = payload_files
|
|
meta["saved_payload_shapes"] = payload_shapes
|
|
meta["saved_payload_dtypes"] = payload_dtypes
|
|
|
|
else:
|
|
if packed_raw is None:
|
|
raise ValueError("packed_raw não pode ser None quando frame_type='RAW_BRUTO'")
|
|
|
|
payload_path = os.path.join(base_dir, f"{name}.bin")
|
|
packed_raw.tofile(payload_path)
|
|
|
|
meta["saved_payload_type"] = "raw_native_single"
|
|
meta["saved_payload_path"] = os.path.basename(payload_path)
|
|
meta["saved_payload_dtype"] = str(packed_raw.dtype)
|
|
meta["saved_payload_shape"] = list(packed_raw.shape)
|
|
|
|
else:
|
|
raise ValueError(f"frame_type não suportado para save: {frame_type}")
|
|
|
|
cv2.imwrite(png_path, preview_bgr)
|
|
|
|
with open(json_path, "w", encoding="utf-8") as f:
|
|
json.dump(meta, f, ensure_ascii=False, indent=2)
|
|
|
|
return png_path, json_path
|
|
|
|
|
|
def get_camera_map_from_status(status: dict) -> dict:
|
|
result = {}
|
|
for cam in status.get("cameras", []):
|
|
result[cam.get("id")] = cam
|
|
return result
|
|
|
|
|
|
# =========================
|
|
# MAIN
|
|
# =========================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Captura de dataset usando módulo multispectral Pi + StreamReceiver.",
|
|
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("--fps", type=int, default=20, help="FPS desejado.")
|
|
parser.add_argument("--width", type=int, default=RAW_SIZE[0], help="Largura óptica da câmera.")
|
|
parser.add_argument("--height", type=int, default=RAW_SIZE[1], help="Altura óptica da câmera.")
|
|
parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
|
|
parser.add_argument("--preview_upscale", type=int, default=2, help="Fator de upscale visual do preview.")
|
|
parser.add_argument("--bayer", default="RGGB", choices=["GBRG", "GRBG", "RGGB", "BGGR"], help="Padrão Bayer das câmeras.")
|
|
parser.add_argument("--output_dtype", default="float32", choices=["uint8", "uint16", "float32"], help="Dtype do payload processado no Pi.")
|
|
parser.add_argument("--frame_type", default="RAW_BRUTO", choices=["RAW_BRUTO", "RGB", "MULTISPEC"], help="Tipo de payload pedido ao Pi.")
|
|
parser.add_argument("--capture_mode", default="AUTO", choices=["AUTO", "SINGLE", "DOUBLE", "TRIPLE"], help="Modo de captura desejado no módulo.")
|
|
parser.add_argument("--raw_policy", default="allow_single", choices=["allow_single", "require_triple"], help="Quando frame_type=RAW_BRUTO, define se o script aceita 1 câmera ou exige 3.")
|
|
parser.add_argument("--module_calibration_json", default=MODULE_PARAMS, help="JSON salvo pelo calibrador de sensores com parâmetros fixos por câmera.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
effective_capture_mode = args.capture_mode
|
|
|
|
raw_w = args.width
|
|
raw_h = args.height
|
|
|
|
session_dir = os.path.join(
|
|
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 - Módulo Multiespectral")
|
|
print(f"Cana : {args.cana}")
|
|
print(f"Horário : {args.horario}")
|
|
print(f"Saída : {session_dir}")
|
|
print(f"Sensor : {raw_w}x{raw_h} | Bayer={args.bayer}")
|
|
print(f"FrameType : {args.frame_type}")
|
|
print(f"CaptureMode : {args.capture_mode} -> efetivo={effective_capture_mode}")
|
|
print(f"RAW policy : {args.raw_policy}")
|
|
print("============================================")
|
|
|
|
beauty_preview = False
|
|
radiometric_ae = True
|
|
auto_save = False
|
|
last_auto_t = 0.0
|
|
preview_upscale = args.preview_upscale
|
|
|
|
t_view_fps = time.time()
|
|
view_frames = 0
|
|
fps_view = 0.0
|
|
|
|
t_stream_fps = time.time()
|
|
last_stream_frame_id = None
|
|
stream_frames_accum = 0
|
|
fps_stream = 0.0
|
|
|
|
last_msg = ""
|
|
last_msg_t = 0.0
|
|
|
|
window_name = "Dataset Capture (C/SPACE=save | A=auto-save | M=preview | R=rad | Q=quit)"
|
|
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
|
|
|
|
last_frame_id = -1
|
|
last_payload_float = None
|
|
last_packed_raw = None
|
|
last_packed_raw_by_camera = None
|
|
last_preview_bgr = None
|
|
last_meta_stream = None
|
|
|
|
try:
|
|
with MultiSpectralClient(
|
|
width=raw_w,
|
|
height=raw_h,
|
|
bayer=args.bayer,
|
|
fps=args.fps,
|
|
frame_type=args.frame_type,
|
|
output_dtype=args.output_dtype,
|
|
capture_mode=effective_capture_mode,
|
|
raw_policy=args.raw_policy,
|
|
module_calibration_json=args.module_calibration_json,
|
|
) as cam:
|
|
while True:
|
|
t0 = time.time()
|
|
frame, meta, decoded = cam.get_next_decoded(timeout=1.0)
|
|
if meta is not None and frame is not None and meta.get("frame_id") != last_frame_id:
|
|
last_frame_id = meta["frame_id"]
|
|
|
|
try:
|
|
frame_type = meta.get("frame_type", "RAW_BRUTO")
|
|
dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8")
|
|
preview_source_id = "rgb"
|
|
|
|
if frame_type == "RAW_BRUTO":
|
|
if isinstance(frame, dict):
|
|
packed_by_camera = frame
|
|
|
|
preview_bgr, raw3_preview, preview_source_id = cam.build_preview_from_raw_payload(frame=frame, meta=meta)
|
|
|
|
if beauty_preview:
|
|
rgb_preview = None
|
|
previews = cam.build_visual_preview_from_raw(frame, meta)
|
|
camera_info = meta.get("camera_info", {}) or {}
|
|
for cam_id, img in previews.items():
|
|
role = camera_info.get(cam_id, {}).get("role")
|
|
if role == "rgb":
|
|
rgb_preview = img
|
|
preview_source_id = cam_id
|
|
break
|
|
if rgb_preview is not None:
|
|
preview_bgr = rgb_preview
|
|
|
|
last_packed_raw = None
|
|
last_packed_raw_by_camera = {cam_id: arr.copy() for cam_id, arr in packed_by_camera.items()}
|
|
last_payload_float = raw3_preview.copy()
|
|
|
|
elif frame_type == "RGB":
|
|
rgb_chw = frame
|
|
if not isinstance(rgb_chw, np.ndarray) or rgb_chw.ndim != 3:
|
|
raise RuntimeError(f"Frame RGB inválido: type={type(rgb_chw)}")
|
|
|
|
if dtype_str == "uint8":
|
|
payload_float = rgb_chw.astype(np.float32) / 255.0
|
|
elif dtype_str == "float32":
|
|
payload_float = rgb_chw.astype(np.float32)
|
|
elif dtype_str == "uint16":
|
|
payload_float = rgb_chw.astype(np.float32) / 65535.0
|
|
else:
|
|
raise RuntimeError(f"dtype RGB não suportado: {dtype_str}")
|
|
|
|
preview_rgb = np.transpose(payload_float, (1, 2, 0))
|
|
preview_bgr = cv2.cvtColor(
|
|
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
|
|
cv2.COLOR_RGB2BGR
|
|
)
|
|
|
|
last_payload_float = payload_float.copy()
|
|
last_packed_raw = None
|
|
last_packed_raw_by_camera = None
|
|
|
|
elif frame_type == "MULTISPEC":
|
|
multispec_chw = frame
|
|
if not isinstance(multispec_chw, np.ndarray) or multispec_chw.ndim != 3 or multispec_chw.shape[0] not in (4, 5):
|
|
raise RuntimeError(f"Frame MULTISPEC inválido: shape={getattr(multispec_chw, 'shape', None)}")
|
|
|
|
if dtype_str == "uint8":
|
|
payload_float = multispec_chw.astype(np.float32) / 255.0
|
|
elif dtype_str == "float32":
|
|
payload_float = multispec_chw.astype(np.float32)
|
|
elif dtype_str == "uint16":
|
|
payload_float = multispec_chw.astype(np.float32) / 65535.0
|
|
else:
|
|
raise RuntimeError(f"dtype MULTISPEC não suportado: {dtype_str}")
|
|
|
|
preview_rgb = np.transpose(payload_float[:3], (1, 2, 0))
|
|
preview_bgr = cv2.cvtColor(
|
|
np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8),
|
|
cv2.COLOR_RGB2BGR
|
|
)
|
|
|
|
last_payload_float = payload_float.copy()
|
|
last_packed_raw = None
|
|
last_packed_raw_by_camera = None
|
|
|
|
else:
|
|
raise RuntimeError(f"frame_type não suportado neste script: {frame_type}")
|
|
|
|
if preview_upscale and preview_upscale > 1:
|
|
preview_show = cv2.resize(
|
|
preview_bgr,
|
|
(preview_bgr.shape[1] * preview_upscale, preview_bgr.shape[0] * preview_upscale),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
else:
|
|
preview_show = preview_bgr.copy()
|
|
|
|
curr_frame_id = meta.get("frame_id")
|
|
|
|
if curr_frame_id is not None:
|
|
if last_stream_frame_id != curr_frame_id:
|
|
stream_frames_accum += 1
|
|
|
|
last_stream_frame_id = curr_frame_id
|
|
|
|
dt_stream = time.time() - t_stream_fps
|
|
if dt_stream >= 1.0:
|
|
fps_stream = stream_frames_accum / dt_stream
|
|
stream_frames_accum = 0
|
|
t_stream_fps = time.time()
|
|
|
|
view_frames += 1
|
|
dt_view = time.time() - t_view_fps
|
|
if dt_view >= 1.0:
|
|
fps_view = view_frames / dt_view
|
|
view_frames = 0
|
|
t_view_fps = time.time()
|
|
|
|
active_sources = meta.get("payload_sources")
|
|
|
|
rad = getattr(cam, "radiometric_controller", None)
|
|
if rad and rad.enabled:
|
|
st = rad.state
|
|
line_rad = (
|
|
f"RAD | "
|
|
f"RGB(exp={st['rgb']['exp']}, g={st['rgb']['gain']:.2f}) | "
|
|
f"RE(exp={st['re']['exp']}, g={st['re']['gain']:.2f}) | "
|
|
f"NIR(exp={st['nir']['exp']}, g={st['nir']['gain']:.2f})"
|
|
)
|
|
else:
|
|
line_rad = "RAD | OFF"
|
|
|
|
lines = [
|
|
f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
|
|
f"Type={meta.get('frame_type')} | CaptureMode={effective_capture_mode} | RAW policy={args.raw_policy}",
|
|
f"Sources={active_sources} | FPS_STREAM={fps_stream:.1f} | FPS_VIEW={fps_view:.1f}",
|
|
f"frame_id={meta.get('frame_id')} | layout={meta.get('output_layout')} | dtype={meta.get('dtype') or meta.get('output_dtype')}",
|
|
f"codec={meta.get('codec_name', meta.get('codec_family', '-'))} | comp={meta.get('dt_comp', 0):.4f}s | send={meta.get('dt_send_payload_prev', 0):.4f}s",
|
|
f"CAM_PARAMS={os.path.basename(args.module_calibration_json)} | controles fixos aplicados",
|
|
line_rad,
|
|
"Keys: C/SPACE=save | A=auto-save | M=preview | R=rad | Q/Esc=quit"
|
|
]
|
|
overlay_hud(preview_show, lines, base_h=raw_h)
|
|
|
|
if last_msg and (time.time() - last_msg_t) < 2.0:
|
|
cv2.putText(preview_show, last_msg, (12, preview_show.shape[0] - 18),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)
|
|
|
|
cv2.imshow(window_name, preview_show)
|
|
|
|
last_preview_bgr = preview_bgr.copy()
|
|
last_meta_stream = dict(meta)
|
|
|
|
except Exception as e:
|
|
err = np.zeros((500, 1200, 3), dtype=np.uint8)
|
|
cv2.putText(err, f"Erro ao processar frame: {e}", (20, 60),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2, cv2.LINE_AA)
|
|
cv2.imshow(window_name, err)
|
|
print(f"[ERRO FRAME] {e}")
|
|
|
|
now = time.time()
|
|
can_save = (
|
|
last_meta_stream is not None and
|
|
last_preview_bgr is not None and
|
|
(
|
|
(last_meta_stream.get("frame_type") in ("RGB", "MULTISPEC") and last_payload_float is not None) or
|
|
(last_meta_stream.get("frame_type") == "RAW_BRUTO" and (last_packed_raw is not None or last_packed_raw_by_camera is not None))
|
|
)
|
|
)
|
|
|
|
if auto_save and can_save and (now - last_auto_t) >= args.interval:
|
|
frame_type_save = last_meta_stream.get("frame_type")
|
|
|
|
meta_save = {
|
|
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
|
"cana": args.cana,
|
|
"horario": args.horario,
|
|
"sensor_width": raw_w,
|
|
"sensor_height": raw_h,
|
|
"bayer_pattern": args.bayer,
|
|
"fps_target": args.fps,
|
|
"frame_type": frame_type_save,
|
|
"capture_mode_requested": args.capture_mode,
|
|
"capture_mode_effective": effective_capture_mode,
|
|
"raw_policy": args.raw_policy,
|
|
"stream_meta": last_meta_stream,
|
|
"startup_camera_controls": cam.applied_camera_controls,
|
|
"actual_camera_controls": cam.get_current_camera_controls(),
|
|
"radiometric_last_result": cam.get_radiometric_last_result(),
|
|
"camera_params_json": args.module_calibration_json,
|
|
"note": "autosave",
|
|
"raw_preview_reference_camera": preview_source_id,
|
|
}
|
|
|
|
save_sample(
|
|
session_dir,
|
|
frame_type=frame_type_save,
|
|
preview_bgr=last_preview_bgr,
|
|
meta=meta_save,
|
|
raw_payload=last_payload_float,
|
|
packed_raw=last_packed_raw,
|
|
packed_raw_by_camera=last_packed_raw_by_camera,
|
|
)
|
|
|
|
last_msg = "SALVO (auto)"
|
|
last_msg_t = now
|
|
last_auto_t = now
|
|
|
|
k = cv2.waitKey(1) & 0xFF
|
|
if k in (ord("q"), ord("Q"), 27):
|
|
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("m"), ord("M")):
|
|
#preview_upscale = 0 if preview_upscale else args.preview_upscale
|
|
beauty_preview = False if beauty_preview else True
|
|
last_msg = f"Preview Beauty -> {beauty_preview}"
|
|
last_msg_t = time.time()
|
|
|
|
elif k in (ord("r"), ord("R")):
|
|
rad = getattr(cam, "radiometric_controller", None)
|
|
if rad is not None:
|
|
radiometric_ae = False if radiometric_ae else True
|
|
rad.enabled = radiometric_ae
|
|
last_msg = f"RAD -> {radiometric_ae}"
|
|
last_msg_t = time.time()
|
|
|
|
|
|
elif k in (ord("c"), ord("C"), 32):
|
|
if can_save:
|
|
frame_type_save = last_meta_stream.get("frame_type")
|
|
meta_save = {
|
|
"ts": datetime.now().isoformat(timespec="milliseconds"),
|
|
"cana": args.cana,
|
|
"horario": args.horario,
|
|
"sensor_width": raw_w,
|
|
"sensor_height": raw_h,
|
|
"bayer_pattern": args.bayer,
|
|
"fps_target": args.fps,
|
|
"frame_type": frame_type_save,
|
|
"capture_mode_requested": args.capture_mode,
|
|
"capture_mode_effective": effective_capture_mode,
|
|
"raw_policy": args.raw_policy,
|
|
"stream_meta": last_meta_stream,
|
|
"startup_camera_controls": cam.applied_camera_controls,
|
|
"actual_camera_controls": cam.get_current_camera_controls(),
|
|
"radiometric_last_result": cam.get_radiometric_last_result(),
|
|
"camera_params_json": args.module_calibration_json,
|
|
"note": "manual",
|
|
"raw_preview_reference_camera": preview_source_id,
|
|
}
|
|
|
|
save_sample(
|
|
session_dir,
|
|
frame_type=frame_type_save,
|
|
preview_bgr=last_preview_bgr,
|
|
meta=meta_save,
|
|
raw_payload=last_payload_float,
|
|
packed_raw=last_packed_raw,
|
|
packed_raw_by_camera=last_packed_raw_by_camera,
|
|
)
|
|
|
|
last_msg = "SALVO (manual)"
|
|
last_msg_t = time.time()
|
|
|
|
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() |