agrobot_base/Python/raspi/cam_2/capture_dataset.py

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,
    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", "RGBNIR"):
        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"] = "raw10_packed_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"] = "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 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


def validate_module_ready(status: dict, frame_type: str, raw_policy: str):
    if not status.get("ok", True):
        raise RuntimeError(f"Status inválido retornado pelo módulo: {status}")

    active_ids = list(status.get("active_camera_ids", []))
    active_count = int(status.get("camera_count_active", 0))
    cam_map = get_camera_map_from_status(status)

    if frame_type == "RGB":
        if "cam0" not in active_ids:
            raise RuntimeError(
                "Modo RGB requer cam0 ativa, mas o módulo não reportou cam0 como ativa."
            )
        return

    if frame_type == "RGBNIR":
        missing = [cid for cid in ("cam0", "cam1") if cid not in active_ids]
        if missing:
            raise RuntimeError(
                f"Modo RGBNIR requer cam0 e cam1 ativas. Faltando: {missing}. "
                f"Ativas atuais: {active_ids}"
            )
        return

    if frame_type == "RAW_BRUTO":
        if raw_policy == "require_dual":
            missing = [cid for cid in ("cam0", "cam1") if cid not in active_ids]
            if missing:
                raise RuntimeError(
                    f"RAW_BRUTO com política require_dual exige duas câmeras ativas. "
                    f"Faltando: {missing}. Ativas atuais: {active_ids}"
                )
        else:
            if active_count < 1:
                raise RuntimeError("RAW_BRUTO requer ao menos uma câmera ativa, mas nenhuma foi detectada.")
        return

    raise RuntimeError(f"frame_type desconhecido para validação: {frame_type}")


def build_preview_from_raw_cam0(
    packed_cam0: np.ndarray,
    processor_core: RawProcessorCore,
    processor_preview: RawProcessorPreview,
    bit_depth: int = 10,
):
    packed = packed_cam0
    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=bit_depth,
    )

    raw3 = processor_core.build_training_rgb(
        raw16,
        output_dtype="float32",
        bit_depth=bit_depth,
    )

    return preview_bgr, raw3, packed


def resolve_effective_capture_mode(frame_type: str, raw_policy: str, requested_mode: str) -> str:
    """
    Decide o modo real que será pedido ao módulo, priorizando segurança e economia.

    Regras:
    - RGB -> SINGLE
    - RGBNIR -> DUAL
    - RAW_BRUTO + allow_single -> SINGLE
    - RAW_BRUTO + require_dual -> DUAL
    """
    if frame_type == "RGB":
        return "SINGLE"

    if frame_type == "RGBNIR":
        return "DUAL"

    if frame_type == "RAW_BRUTO":
        if raw_policy == "require_dual":
            return "DUAL"
        return "SINGLE"

    return requested_mode

# =========================
# 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("--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 das câmeras.")
    parser.add_argument("--frame_type", default="RAW_BRUTO", choices=["RAW_BRUTO", "RGB", "RGBNIR"], help="Tipo de payload pedido ao Pi.")
    parser.add_argument("--output_dtype", default="float32", choices=["uint8", "uint16", "float32"], help="Dtype do payload processado no Pi.")
    parser.add_argument("--capture_mode", default="AUTO", choices=["AUTO", "SINGLE", "DUAL"], help="Modo de captura desejado no módulo.")
    parser.add_argument("--raw_policy", default="allow_single", choices=["allow_single", "require_dual"], help="Quando frame_type=RAW_BRUTO, define se o script aceita 1 câmera ou exige 2.")

    args = parser.parse_args()

    effective_capture_mode = resolve_effective_capture_mode(
        frame_type=args.frame_type,
        raw_policy=args.raw_policy,
        requested_mode=args.capture_mode,
    )

    raw_w = args.width
    raw_h = args.height

    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 - 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("============================================")

    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)

    print(f"[INFO] Verificando conexão com o módulo em {args.pi_host}:{args.server_port}...")

    svc.ensure_alive()

    if not svc.is_alive():
        raise RuntimeError("Módulo não encontrado ou não respondeu ao ping.")

    print("[OK] Módulo conectado e respondendo.")

    window_name = "Dataset Capture (C/SPACE=save | A=auto-save | M=preview scale | Q=quit)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    processor_core_cam0 = RawProcessorCore(
        sensor_width=raw_w,
        sensor_height=raw_h,
        bayer_pattern=args.bayer,
    )
    processor_preview_cam0 = RawProcessorPreview(
        sensor_width=raw_w,
        sensor_height=raw_h,
        bayer_pattern=args.bayer,
    )

    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:
        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 CAM0 RES:", svc.set_camera_resolution(0, raw_w, raw_h))
        print("SET CAM1 RES:", svc.set_camera_resolution(1, raw_w, raw_h))
        print("SET CAM0 BAYER:", svc.set_camera_bayer(0, args.bayer))
        print("SET CAM1 BAYER:", svc.set_camera_bayer(1, args.bayer))
        print("SET FPS:", svc.set_fps(args.fps))
        print("SET CAPTURE MODE:", svc.set_capture_mode(effective_capture_mode))
        print("SET FRAME TYPE:", svc.set_frame_type(args.frame_type))
        print("SET OUTPUT DTYPE:", svc.set_output_dtype(args.output_dtype))

        begin_resp = svc.begin(
            frame_type=args.frame_type,
            output_dtype=args.output_dtype,
            capture_mode=effective_capture_mode,
        )
        print("BEGIN:", begin_resp)

        status = svc.get_status()
        print("STATUS:", json.dumps({
            "status": status.get("status"),
            "detected_mode": status.get("detected_mode"),
            "camera_count_active": status.get("camera_count_active"),
            "active_camera_ids": status.get("active_camera_ids"),
        }, ensure_ascii=False))

        validate_module_ready(status, args.frame_type, args.raw_policy)

        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)

        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"]

                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":
                        if isinstance(frame, dict):
                            packed_by_camera = frame
                            cam0 = packed_by_camera.get("cam0")
                            if cam0 is None:
                                raise RuntimeError("RAW_BRUTO multi recebido sem cam0 para preview")

                            source_cam0 = None
                            for src in (meta.get("raw_sources") or []):
                                if src.get("id") == "cam0":
                                    source_cam0 = src
                                    break

                            bit_depth_cam0 = 10 if source_cam0 is None else int(source_cam0.get("bit_depth", 10))

                            preview_bgr, raw3_preview, packed_cam0 = build_preview_from_raw_cam0(
                                cam0,
                                processor_core_cam0,
                                processor_preview_cam0,
                                bit_depth=bit_depth_cam0,
                            )

                            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()

                        else:
                            source_camera = meta.get("source_camera") or {}
                            bit_depth = int(source_camera.get("bit_depth", meta.get("source_bit_depth", 10)))

                            preview_bgr, raw3_preview, packed_single = build_preview_from_raw_cam0(
                                frame,
                                processor_core_cam0,
                                processor_preview_cam0,
                                bit_depth=bit_depth,
                            )

                            last_packed_raw = packed_single.copy()
                            last_packed_raw_by_camera = None
                            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 == "RGBNIR":
                        rgbnir_chw = frame
                        if not isinstance(rgbnir_chw, np.ndarray) or rgbnir_chw.ndim != 3 or rgbnir_chw.shape[0] != 5:
                            raise RuntimeError(f"Frame RGBNIR inválido: shape={getattr(rgbnir_chw, 'shape', None)}")

                        if dtype_str == "uint8":
                            payload_float = rgbnir_chw.astype(np.float32) / 255.0
                        elif dtype_str == "float32":
                            payload_float = rgbnir_chw.astype(np.float32)
                        elif dtype_str == "uint16":
                            payload_float = rgbnir_chw.astype(np.float32) / 65535.0
                        else:
                            raise RuntimeError(f"dtype RGBNIR 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 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()

                    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")
                    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"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_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", "RGBNIR") 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,
                    "note": "autosave",
                }

                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
                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")):
                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")):
                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")):
                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")):
                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")):
                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):
                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": args.capture_mode,
                        "raw_policy": args.raw_policy,
                        "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",
                    }

                    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:
        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()