import socket import json import numpy as np import base64 import time from typing import Optional class MultiSpectralService: def __init__(self, host="192.168.105.6", port=5000, timeout=5): self.host = host self.port = port self.timeout = timeout self.sock = None self.file = None def __enter__(self): self.connect() return self def __exit__(self, exc_type, exc, tb): self.disconnect() def connect(self): if self.sock is not None: return self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout) self.sock.settimeout(self.timeout) self.file = self.sock.makefile("r", encoding="utf-8") def disconnect(self): try: if self.file: self.file.close() except: pass try: if self.sock: self.sock.close() except: pass self.file = None self.sock = None def _send_command(self, payload: dict) -> dict: if self.sock is None: self.connect() data = (json.dumps(payload) + "\n").encode("utf-8") self.sock.sendall(data) line = self.file.readline() if not line: self.disconnect() raise RuntimeError("Conexão encerrada pelo servidor") return json.loads(line.strip()) def _numpy_dtype_from_string(self, dtype_str: str): mapping = { "uint8": np.uint8, "float32": np.float32, "uint16": np.uint16, } if dtype_str not in mapping: raise RuntimeError(f"dtype não suportado recebido do Pi: {dtype_str}") return mapping[dtype_str] def ping(self): return self._send_command({"cmd": "ping"}) def get_status(self): return self._send_command({"cmd": "get_status"}) def get_config(self): return self._send_command({"cmd": "get_config"}) def begin(self, frame_type: str = "RAW_BRUTO", output_dtype: str = "uint8"): return self._send_command({ "cmd": "begin", "frame_type": frame_type, "output_dtype": output_dtype, }) def stop(self): return self._send_command({"cmd": "stop"}) def set_fps(self, fps: int): return self._send_command({"cmd": "set_fps", "value": fps}) def set_jpeg_quality(self, quality: int): return self._send_command({"cmd": "set_jpeg_quality", "value": quality}) def set_bayer(self, bayer_pattern: str): return self._send_command({ "cmd": "set_bayer", "pattern": bayer_pattern }) def set_resolution(self, width: int, height: int): return self._send_command({ "cmd": "set_resolution", "width": width, "height": height }) def capture_frame_array(self): t0 = time.perf_counter() resp = self._send_command({"cmd": "capture_frame"}) if not resp.get("ok"): raise RuntimeError(resp.get("error", "Falha ao capturar frame")) raw = base64.b64decode(resp["data"]) frame_type = resp.get("frame_type") output_layout = resp.get("output_layout", "HWC") dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8" np_dtype = self._numpy_dtype_from_string(dtype_str) width = int(resp.get("output_width", resp.get("width"))) height = int(resp.get("output_height", resp.get("height"))) channels = int(resp.get("output_channels", resp.get("channels", 1))) arr = np.frombuffer(raw, dtype=np_dtype) if output_layout == "HW": arr = arr.reshape(height, width) elif output_layout == "CHW": arr = arr.reshape(channels, height, width) elif output_layout == "HWC": arr = arr.reshape(height, width, channels) else: raise RuntimeError(f"Layout não suportado recebido do Pi: {output_layout}") t1 = time.perf_counter() meta = { "frame_type": frame_type, "payload_format_version": resp.get("payload_format_version"), "width": width, "height": height, "channels": channels, "dtype": dtype_str, "output_dtype": resp.get("output_dtype"), "output_layout": output_layout, "output_channel_names": resp.get("output_channel_names"), "packed_width": resp.get("packed_width"), "packed_height": resp.get("packed_height"), "source_width": resp.get("source_width"), "source_height": resp.get("source_height"), "source_bayer_pattern": resp.get("source_bayer_pattern"), "source_bit_depth": resp.get("source_bit_depth"), "size": resp.get("size"), "ts_pi": resp.get("ts_pi"), "ts_pi_monotonic": resp.get("ts_pi_monotonic"), "dt_trigger": resp.get("dt_trigger"), "dt_settle": resp.get("dt_settle"), "dt_capture": resp.get("dt_capture"), "dt_process": resp.get("dt_process"), "dt_total_pi": resp.get("dt_total_pi"), "dt_total_pc": t1 - t0, } return arr, meta def start_stream(self, host: str, port: int, fps: float): return self._send_command({ "cmd": "start_stream", "host": host, "port": port, "fps": fps }) def stop_stream(self): return self._send_command({"cmd": "stop_stream"}) def get_camera_controls(self): return self._send_command({"cmd": "get_camera_controls"}) def set_ae_enable(self, value: bool): return self._send_command({"cmd": "set_ae_enable", "value": bool(value)}) def set_awb_enable(self, value: bool): return self._send_command({"cmd": "set_awb_enable", "value": bool(value)}) def set_exposure_time(self, exposure_time_us: Optional[int] = None): return self._send_command({"cmd": "set_exposure_time", "value": exposure_time_us}) def clear_exposure_time(self): return self._send_command({"cmd": "clear_exposure_time"}) def set_analogue_gain(self, gain: float | None): return self._send_command({"cmd": "set_analogue_gain", "value": gain}) def clear_analogue_gain(self): return self._send_command({"cmd": "clear_analogue_gain"}) def set_colour_gains(self, r_gain: float, b_gain: float): return self._send_command({ "cmd": "set_colour_gains", "r_gain": r_gain, "b_gain": b_gain }) def clear_colour_gains(self): return self._send_command({"cmd": "clear_colour_gains"}) def get_sensor_modes(self): return self._send_command({"cmd": "get_sensor_modes"})