import os import time import json import argparse from datetime import datetime import cv2 import numpy as np from multispectral_service import MultiSpectralService from stream_receiver import StreamReceiver from pi.raw_processor_core import RawProcessorCore from pi.raw_processor_preview import RawProcessorPreview STREAM_PORT = 6001 PI_HOST = "192.168.105.6" PC_HOST = "192.168.105.5" # ========================= # 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, raw3: np.ndarray | None = None, packed_raw: np.ndarray | None = None, ): """ Salva conforme o tipo de frame: - RGB: * payload em .raw float32 (3,H,W) * preview em .png * metadados em .json - RAW_BRUTO: * payload packed RAW10 em .bin * preview em .png * metadados em .json """ 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 == "RGB": if raw3 is None: raise ValueError("raw3 não pode ser None quando frame_type='RGB'") payload_path = os.path.join(base_dir, f"{name}.raw") raw3.astype(np.float32).tofile(payload_path) meta["saved_payload_type"] = "raw3" meta["saved_payload_path"] = os.path.basename(payload_path) meta["saved_payload_dtype"] = "float32" meta["saved_payload_shape"] = list(raw3.shape) elif frame_type == "RAW_BRUTO": 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"] = "raw10_packed" 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 payload_path, png_path, json_path # ========================= # MAIN # ========================= def main(): parser = argparse.ArgumentParser( description="Captura de dataset RAW3 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("--modelo", default="imx296_pi", help="Nome do módulo/câmera para montar a pasta.") parser.add_argument("--pi_host", default=PI_HOST, help="IP do servidor no Raspberry Pi.") parser.add_argument("--pc_host", default=PC_HOST, help="IP local do notebook/PC que receberá o stream.") parser.add_argument("--stream_port", type=int, default=STREAM_PORT, help="Porta TCP do receiver de stream.") parser.add_argument("--server_port", type=int, default=5000, help="Porta TCP do servidor de comandos no Pi.") parser.add_argument("--fps", type=int, default=20, help="FPS desejado.") parser.add_argument("--width", type=int, default=640, help="Largura óptica da câmera.") parser.add_argument("--height", type=int, default=480, 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="GBRG", choices=["GBRG", "GRBG", "RGGB", "BGGR"], help="Padrão Bayer da câmera.") parser.add_argument("--codec_family", default="numcodecs", help="Apenas informativo no metadata local.") parser.add_argument("--codec_name", default="blosc", help="Apenas informativo no metadata local.") parser.add_argument("--frame_type", default="RAW_BRUTO", choices=["RAW_BRUTO", "RGB"], help="Tipo de payload pedido ao Pi.") parser.add_argument("--output_dtype", default="uint8", choices=["uint8", "float32"], help="Dtype do payload processado no Pi.") args = parser.parse_args() raw_w = args.width raw_h = args.height # dataset//brutas/cana_/// session_dir = os.path.join( args.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 RAW3 - 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("============================================") window_name = "Dataset Capture - RAW3 (C/SPACE=save | A=auto-save | M=preview scale | Q=quit)" cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) 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 receiver = StreamReceiver(host="0.0.0.0", port=args.stream_port) svc = MultiSpectralService(host=args.pi_host, port=args.server_port, timeout=10) processor_core = RawProcessorCore( sensor_width=raw_w, sensor_height=raw_h, bayer_pattern=args.bayer, ) processor_preview = RawProcessorPreview( sensor_width=raw_w, sensor_height=raw_h, bayer_pattern=args.bayer, ) last_frame_id = -1 last_raw3 = None last_packed_raw = None last_preview_bgr = None last_meta_stream = None try: receiver.start() time.sleep(0.5) svc.connect() modes_resp = svc.get_sensor_modes() if not modes_resp.get("ok"): raise RuntimeError("Falha ao obter sensor_modes") for mode in modes_resp["sensor_modes"]: print( f"[{mode['index']}] " f"size={mode['size']} " f"format={mode['format']} " f"bit_depth={mode['bit_depth']} " f"fps={mode['fps']}" ) print("SET RES:", svc.set_resolution(raw_w, raw_h)) print("SET FPS:", svc.set_fps(args.fps)) print("SET BAYER:", svc.set_bayer(args.bayer)) print("BEGIN:", svc.begin(frame_type=args.frame_type, output_dtype=args.output_dtype)) print("START STREAM:", svc.start_stream(args.pc_host, args.stream_port, fps=args.fps)) camera_ctrl = svc.get_camera_controls() ae_enabled = bool(camera_ctrl.get("ae_enable", True)) awb_enabled = bool(camera_ctrl.get("awb_enable", True)) manual_exposure_us = camera_ctrl.get("exposure_time_us", None) manual_gain = camera_ctrl.get("analogue_gain", None) manual_colour_gains = camera_ctrl.get("colour_gains", None) expected_packed_w = (args.width * 10) // 8 expected_packed_h = args.height while True: t0 = time.time() meta = receiver.last_meta frame = receiver.last_frame if meta is not None and frame is not None and meta.get("frame_id") != last_frame_id: last_frame_id = meta["frame_id"] #print(meta) if False and meta["packed_width"] != expected_packed_w or meta["packed_height"] != expected_packed_h: def infer_sensor_size_from_packed(packed_width: int, packed_height: int): width = int((packed_width * 8) / 10) height = packed_height return width, height actual_packed_w = meta["packed_width"] actual_packed_h = meta["packed_height"] padding = actual_packed_w - expected_packed_w if actual_packed_h != args.height: erro = True elif actual_packed_w < expected_packed_w: erro = True elif padding > 64: # margem conservadora, se quiser erro = True else: erro = False if erro: inferred_w = int((actual_packed_w * 8) / 10) inferred_h = actual_packed_h modes_txt = [] if modes_resp.get("ok"): for m in modes_resp["sensor_modes"]: size = m.get("size") fmt = m.get("format") fps = m.get("fps") modes_txt.append(f"- {size[0]}x{size[1]} | {fmt} | fps={fps}") modes_str = "\n".join(modes_txt) if modes_txt else "(não disponível)" RuntimeError( f"Modo RAW inesperado.\n" f"Solicitado: {args.width}x{args.height} (packed útil esperado {expected_packed_w})\n" f"Recebido: packed {actual_packed_h}x{actual_packed_w}\n" f"Obs: packed_width pode incluir padding/stride.\n" f"Se a altura confere e o packed recebido é maior que o esperado, " f"o modo pode estar correto com alinhamento de memória." ) try: frame_type = meta.get("frame_type", "RAW_BRUTO") dtype_str = meta.get("dtype") or meta.get("output_dtype", "uint8") if frame_type == "RAW_BRUTO": packed = frame if packed.ndim == 3 and packed.shape[2] == 1: packed = packed[:, :, 0] raw16 = processor_core.unpack_raw10_packed(packed) preview_bgr = processor_preview.raw16_to_preview_bgr( raw16, bit_depth=meta.get("source_bit_depth", 10), ) # continua gerando raw3 apenas para visualização/depuração local, se quiser manter raw3 = processor_core.build_training_rgb( raw16, output_dtype="float32", bit_depth=meta.get("source_bit_depth", 10), ) last_packed_raw = packed.copy() elif frame_type == "RGB": rgb_chw = frame if rgb_chw.ndim != 3: raise RuntimeError(f"Frame RGB inválido: shape={rgb_chw.shape}") if dtype_str == "uint8": raw3 = rgb_chw.astype(np.float32) / 255.0 elif dtype_str == "float32": raw3 = rgb_chw.astype(np.float32) else: raise RuntimeError(f"dtype RGB não suportado: {dtype_str}") preview_rgb = np.transpose(raw3, (1, 2, 0)) preview_bgr = cv2.cvtColor( np.clip(preview_rgb * 255.0, 0, 255).astype(np.uint8), cv2.COLOR_RGB2BGR ) last_packed_raw = 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() # ========================= # FPS do stream (frames recebidos) # ========================= curr_frame_id = meta.get("frame_id") if last_stream_frame_id is not None and curr_frame_id is not None: delta_ids = curr_frame_id - last_stream_frame_id if delta_ids > 0: stream_frames_accum += delta_ids 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() # ========================= # FPS de visualização/processamento no PC # ========================= 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() lines = [ f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}", f"AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s | PreviewScale: {preview_upscale}", f"frame_id={meta.get('frame_id')} | FPS_STREAM={fps_stream:.1f} | FPS_VIEW={fps_view:.1f}", 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"packed={meta.get('width')}x{meta.get('height')} | raw3_shape={list(raw3.shape)}", f"type={meta.get('frame_type')} | layout={meta.get('output_layout')} | dtype={meta.get('dtype') or meta.get('output_dtype')}", f"AE={'ON' if ae_enabled else 'OFF'} | AWB={'ON' if awb_enabled else 'OFF'} | EXP={manual_exposure_us} | GAIN={manual_gain}", "Keys: C/SPACE=save | A=auto-save | E=AE | W=AWB | I/K=exp | O/L=gain | R=reset | M=preview | 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_raw3 = raw3.copy() if raw3 is not None else None last_preview_bgr = preview_bgr.copy() if preview_bgr is not None else None 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") == "RGB" and last_raw3 is not None) or (last_meta_stream.get("frame_type") == "RAW_BRUTO" and last_packed_raw 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": last_meta_stream.get("source_bayer_pattern"), "fps_target": args.fps, "frame_type": frame_type_save, "codec_family": last_meta_stream.get("codec_family"), "codec_name": last_meta_stream.get("codec_name"), "codec_params": last_meta_stream.get("codec_params"), "stream_meta": last_meta_stream, "note": "autosave", } if frame_type_save == "RGB": meta_save["raw3_shape"] = list(last_raw3.shape) meta_save["raw3_dtype"] = str(last_raw3.dtype) elif frame_type_save == "RAW_BRUTO": meta_save["packed_shape"] = list(last_packed_raw.shape) meta_save["packed_dtype"] = str(last_packed_raw.dtype) meta_save["packed_height"] = int(last_packed_raw.shape[0]) meta_save["packed_width"] = int(last_packed_raw.shape[1]) payload_path, _, _ = save_sample( session_dir, frame_type=frame_type_save, preview_bgr=last_preview_bgr, meta=meta_save, raw3=last_raw3, packed_raw=last_packed_raw, ) last_msg = f"SALVO (auto): {os.path.basename(payload_path)}" 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 last_msg = f"Preview UPSCALE -> {preview_upscale}" last_msg_t = time.time() elif k in (ord("e"), ord("E")): ae_enabled = not ae_enabled resp = svc.set_ae_enable(ae_enabled) ae_enabled = bool(resp.get("ae_enable", ae_enabled)) last_msg = f"AE -> {'ON' if ae_enabled else 'OFF'}" last_msg_t = time.time() elif k in (ord("w"), ord("W")): awb_enabled = not awb_enabled resp = svc.set_awb_enable(awb_enabled) awb_enabled = bool(resp.get("awb_enable", awb_enabled)) last_msg = f"AWB -> {'ON' if awb_enabled else 'OFF'}" last_msg_t = time.time() elif k in (ord("i"), ord("I")): # aumenta exposição manual if manual_exposure_us is None: manual_exposure_us = 15000 else: manual_exposure_us = min(int(manual_exposure_us * 1.15), 200000) if ae_enabled: ae_enabled = False svc.set_ae_enable(False) resp = svc.set_exposure_time(int(manual_exposure_us)) manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us) last_msg = f"ExposureTime -> {manual_exposure_us} us" last_msg_t = time.time() elif k in (ord("k"), ord("K")): # diminui exposição manual if manual_exposure_us is None: manual_exposure_us = 15000 else: manual_exposure_us = max(int(manual_exposure_us / 1.15), 100) if ae_enabled: ae_enabled = False svc.set_ae_enable(False) resp = svc.set_exposure_time(int(manual_exposure_us)) manual_exposure_us = resp.get("exposure_time_us", manual_exposure_us) last_msg = f"ExposureTime -> {manual_exposure_us} us" last_msg_t = time.time() elif k in (ord("o"), ord("O")): # aumenta ganho manual if manual_gain is None: manual_gain = 1.0 else: manual_gain = min(float(manual_gain) * 1.10, 32.0) if ae_enabled: ae_enabled = False svc.set_ae_enable(False) resp = svc.set_analogue_gain(float(manual_gain)) manual_gain = resp.get("analogue_gain", manual_gain) last_msg = f"AnalogueGain -> {manual_gain:.2f}" last_msg_t = time.time() elif k in (ord("l"), ord("L")): # diminui ganho manual if manual_gain is None: manual_gain = 1.0 else: manual_gain = max(float(manual_gain) / 1.10, 1.0) if ae_enabled: ae_enabled = False svc.set_ae_enable(False) resp = svc.set_analogue_gain(float(manual_gain)) manual_gain = resp.get("analogue_gain", manual_gain) last_msg = f"AnalogueGain -> {manual_gain:.2f}" last_msg_t = time.time() elif k in (ord("r"), ord("R")): # reset manuais svc.clear_exposure_time() svc.clear_analogue_gain() svc.clear_colour_gains() manual_exposure_us = None manual_gain = None manual_colour_gains = None last_msg = "Manual controls resetados" last_msg_t = time.time() elif k in (ord("c"), ord("C"), 32): can_save = ( last_meta_stream is not None and last_preview_bgr is not None and ( (last_meta_stream.get("frame_type") == "RGB" and last_raw3 is not None) or (last_meta_stream.get("frame_type") == "RAW_BRUTO" and last_packed_raw is not None) ) ) 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": last_meta_stream.get("source_bayer_pattern"), "fps_target": args.fps, "frame_type": frame_type_save, "codec_family": last_meta_stream.get("codec_family"), "codec_name": last_meta_stream.get("codec_name"), "codec_params": last_meta_stream.get("codec_params"), "stream_meta": last_meta_stream, "camera_controls": { "ae_enable": ae_enabled, "awb_enable": awb_enabled, "exposure_time_us": manual_exposure_us, "analogue_gain": manual_gain, "colour_gains": manual_colour_gains, }, "note": "manual", } if frame_type_save == "RGB": meta_save["raw3_shape"] = list(last_raw3.shape) meta_save["raw3_dtype"] = str(last_raw3.dtype) elif frame_type_save == "RAW_BRUTO": meta_save["packed_shape"] = list(last_packed_raw.shape) meta_save["packed_dtype"] = str(last_packed_raw.dtype) meta_save["packed_height"] = int(last_packed_raw.shape[0]) meta_save["packed_width"] = int(last_packed_raw.shape[1]) payload_path, _, _ = save_sample( session_dir, frame_type=frame_type_save, preview_bgr=last_preview_bgr, meta=meta_save, raw3=last_raw3, packed_raw=last_packed_raw, ) last_msg = f"SALVO (manual): {os.path.basename(payload_path)}" last_msg_t = time.time() dt_loop = time.time() - t0 if dt_loop < 0.001: time.sleep(0.001) finally: try: print("STOP STREAM:", svc.stop_stream()) except Exception: pass try: print("STOP:", svc.stop()) except Exception: pass svc.disconnect() receiver.stop() cv2.destroyAllWindows() print("Fim da captura.") if __name__ == "__main__": main()