import json import socket import threading import time from numcodecs import Blosc import numpy as np class StreamReceiver: def __init__(self, host="0.0.0.0", port=6001): self.host = host self.port = port self._server_sock = None self._client_sock = None self._thread = None self._running = False self.last_frame = None self.last_meta = None self.last_receive_ts = None self._codec = None self._codec_signature = None @property def is_running(self): return self._running def start(self): if self._running: return self._running = True self._thread = threading.Thread(target=self._worker, daemon=True) self._thread.start() def stop(self): self._running = False try: if self._client_sock: self._client_sock.close() except Exception: pass try: if self._server_sock: self._server_sock.close() except Exception: pass self._client_sock = None self._server_sock = None def _recv_exact(self, sock: socket.socket, n: int) -> bytes: chunks = [] remaining = n while remaining > 0: chunk = sock.recv(remaining) if not chunk: raise ConnectionError("Conexão encerrada durante recv") chunks.append(chunk) remaining -= len(chunk) return b"".join(chunks) def _worker(self): try: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server: server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) server.bind((self.host, self.port)) server.listen(1) server.settimeout(1.0) self._server_sock = server print(f"[INFO] StreamReceiver ouvindo em {self.host}:{self.port}") while self._running: try: client, addr = server.accept() except socket.timeout: continue print(f"[INFO] StreamReceiver conectado por {addr}") self._client_sock = client with client: while self._running: header_len = int.from_bytes(self._recv_exact(client, 4), "big") header_bytes = self._recv_exact(client, header_len) header = json.loads(header_bytes.decode("utf-8")) payload_len = int.from_bytes(self._recv_exact(client, 4), "big") payload_comp = self._recv_exact(client, payload_len) self._ensure_codec(header) if header.get("codec_family") == "none": payload = payload_comp else: payload = self._codec.decode(payload_comp) expected = int(header["payload_size_raw"]) if len(payload) != expected: raise ValueError( f"Tamanho descomprimido inválido: {len(payload)} != {expected}" ) if header.get("multi_payload", False): frame = self._decode_multi_payload(payload, header) else: frame = self._decode_single_payload(payload, header) self.last_frame = frame self.last_meta = header self.last_receive_ts = time.perf_counter() print("[INFO] StreamReceiver cliente desconectado") self._client_sock = None except Exception as e: print(f"[WARN] StreamReceiver encerrado com erro: {e}") finally: self._running = False self._client_sock = None self._server_sock = None def _decode_single_payload(self, payload: bytes, header: dict): height = int(header.get("output_height", header.get("height"))) width = int(header.get("output_width", header.get("width"))) channels = int(header.get("output_channels", header.get("channels", 1))) layout = header.get("output_layout", "HWC") dtype = self._numpy_dtype_from_header(header) arr = np.frombuffer(payload, dtype=dtype) if layout == "HW": return arr.reshape(height, width) if layout == "CHW": return arr.reshape(channels, height, width) if layout == "HWC": return arr.reshape(height, width, channels) raise RuntimeError(f"Layout não suportado: {layout}") def _decode_multi_payload(self, payload: bytes, header: dict): payload_parts = header.get("payload_parts", []) camera_frames = header.get("camera_frames", {}) dtype = self._numpy_dtype_from_header(header) frames = {} offset = 0 for part in payload_parts: cam_id = part.get("camera_id") if not cam_id: raise RuntimeError("payload_parts sem camera_id") if offset + 4 > len(payload): raise RuntimeError("Payload multi truncado ao ler tamanho da parte") part_size = int.from_bytes(payload[offset:offset + 4], "big") offset += 4 if offset + part_size > len(payload): raise RuntimeError(f"Payload multi truncado ao ler dados de {cam_id}") part_bytes = payload[offset:offset + part_size] offset += part_size cam_meta = camera_frames.get(cam_id, {}) width = int(cam_meta.get("width")) height = int(cam_meta.get("height")) channels = int(cam_meta.get("channels", 1)) # Para RAW dual atual, cada parte é HW. arr = np.frombuffer(part_bytes, dtype=dtype) if channels == 1: frame = arr.reshape(height, width) else: frame = arr.reshape(channels, height, width) frames[cam_id] = frame if offset != len(payload): raise RuntimeError( f"Payload multi com bytes sobrando: consumidos={offset}, total={len(payload)}" ) return frames def _normalize_shuffle(self, shuffle_value): if isinstance(shuffle_value, int): return shuffle_value mapping = { "NOSHUFFLE": Blosc.NOSHUFFLE, "SHUFFLE": Blosc.SHUFFLE, "BITSHUFFLE": Blosc.BITSHUFFLE, } key = str(shuffle_value).upper() if key not in mapping: raise ValueError(f"shuffle inválido: {shuffle_value}") return mapping[key] def _build_codec_from_header(self, header: dict): family = header.get("codec_family") name = header.get("codec_name") params = dict(header.get("codec_params", {})) if family != "numcodecs": raise ValueError(f"Família de codec não suportada: {family}") if name == "blosc": params["shuffle"] = self._normalize_shuffle(params.get("shuffle", "SHUFFLE")) return Blosc(**params) raise ValueError(f"Codec numcodecs não suportado: {name}") def _get_codec_signature_from_header(self, header: dict): return ( header.get("codec_family"), header.get("codec_name"), tuple(sorted(dict(header.get("codec_params", {})).items())) ) def _ensure_codec(self, header: dict): family = header.get("codec_family") if family == "none": self._codec = None self._codec_signature = ("none", None, ()) return sig = self._get_codec_signature_from_header(header) if self._codec is None or self._codec_signature != sig: self._codec = self._build_codec_from_header(header) self._codec_signature = sig def _numpy_dtype_from_header(self, header: dict): dtype_str = header.get("dtype") or header.get("output_dtype") or "uint8" mapping = { "uint8": np.uint8, "float32": np.float32, "uint16": np.uint16, } if dtype_str == "multi": # fallback atual para RAW bruto multi return np.uint16 if dtype_str not in mapping: raise RuntimeError(f"dtype não suportado: {dtype_str}") return mapping[dtype_str]