324 lines
11 KiB
Python
324 lines
11 KiB
Python
|
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
|
||
|
|
|
||
|
|
# =========================================================
|
||
|
|
# Lifecycle
|
||
|
|
# =========================================================
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
# =========================================================
|
||
|
|
# Socket helpers
|
||
|
|
# =========================================================
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
# =========================================================
|
||
|
|
# Worker principal
|
||
|
|
# =========================================================
|
||
|
|
|
||
|
|
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}"
|
||
|
|
)
|
||
|
|
|
||
|
|
payload_kind = header.get("payload_kind")
|
||
|
|
if header.get("multi_payload", False) or payload_kind == "multi_array":
|
||
|
|
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
|
||
|
|
|
||
|
|
# =========================================================
|
||
|
|
# Decode helpers
|
||
|
|
# =========================================================
|
||
|
|
|
||
|
|
def _numpy_dtype_from_string(self, dtype_str: str):
|
||
|
|
mapping = {
|
||
|
|
"uint8": np.uint8,
|
||
|
|
"uint16": np.uint16,
|
||
|
|
"float32": np.float32,
|
||
|
|
}
|
||
|
|
|
||
|
|
if dtype_str not in mapping:
|
||
|
|
raise RuntimeError(f"dtype não suportado: {dtype_str}")
|
||
|
|
|
||
|
|
return mapping[dtype_str]
|
||
|
|
|
||
|
|
def _numpy_dtype_from_header(self, header: dict):
|
||
|
|
dtype_str = header.get("dtype") or header.get("output_dtype") or "uint8"
|
||
|
|
|
||
|
|
if dtype_str == "multi":
|
||
|
|
# fallback conservador para protocolos mais antigos
|
||
|
|
return np.uint16
|
||
|
|
|
||
|
|
return self._numpy_dtype_from_string(dtype_str)
|
||
|
|
|
||
|
|
def _reshape_from_shape(self, payload: bytes, dtype_str: str, shape):
|
||
|
|
dtype = self._numpy_dtype_from_string(dtype_str)
|
||
|
|
arr = np.frombuffer(payload, dtype=dtype)
|
||
|
|
return arr.reshape(tuple(shape))
|
||
|
|
|
||
|
|
def _reshape_from_layout(self, payload: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
|
||
|
|
dtype = self._numpy_dtype_from_string(dtype_str)
|
||
|
|
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_single_payload(self, payload: bytes, header: dict):
|
||
|
|
payload_parts = header.get("payload_parts", []) or []
|
||
|
|
first_part = payload_parts[0] if payload_parts else {}
|
||
|
|
|
||
|
|
dtype_str = first_part.get("dtype") or header.get("dtype") or header.get("output_dtype") or "uint8"
|
||
|
|
shape = first_part.get("shape")
|
||
|
|
|
||
|
|
if shape:
|
||
|
|
return self._reshape_from_shape(payload, dtype_str, shape)
|
||
|
|
|
||
|
|
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")
|
||
|
|
|
||
|
|
return self._reshape_from_layout(
|
||
|
|
payload=payload,
|
||
|
|
dtype_str=dtype_str,
|
||
|
|
layout=layout,
|
||
|
|
width=width,
|
||
|
|
height=height,
|
||
|
|
channels=channels,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _decode_multi_payload(self, payload: bytes, header: dict):
|
||
|
|
payload_parts = header.get("payload_parts", []) or []
|
||
|
|
camera_frames = header.get("camera_frames", {}) or {}
|
||
|
|
|
||
|
|
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, {})
|
||
|
|
|
||
|
|
dtype_str = part.get("dtype") or header.get("dtype") or header.get("output_dtype")
|
||
|
|
shape = part.get("shape")
|
||
|
|
|
||
|
|
if dtype_str == "multi" or dtype_str is None:
|
||
|
|
channels = int(part.get("channels", cam_meta.get("channels", 1)))
|
||
|
|
dtype_str = "uint8" if channels > 1 else "uint16"
|
||
|
|
|
||
|
|
if shape:
|
||
|
|
frame = self._reshape_from_shape(part_bytes, dtype_str, shape)
|
||
|
|
else:
|
||
|
|
width = int(part.get("width", cam_meta.get("width", 0)))
|
||
|
|
height = int(part.get("height", cam_meta.get("height", 0)))
|
||
|
|
channels = int(part.get("channels", cam_meta.get("channels", 1)))
|
||
|
|
|
||
|
|
layout = "HWC" if channels > 1 else "HW"
|
||
|
|
|
||
|
|
frame = self._reshape_from_layout(
|
||
|
|
payload=part_bytes,
|
||
|
|
dtype_str=dtype_str,
|
||
|
|
layout=layout,
|
||
|
|
width=width,
|
||
|
|
height=height,
|
||
|
|
channels=channels,
|
||
|
|
)
|
||
|
|
|
||
|
|
frames[cam_id] = frame
|
||
|
|
|
||
|
|
if offset != len(payload):
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Payload multi com bytes sobrando: consumidos={offset}, total={len(payload)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
return frames
|
||
|
|
|
||
|
|
# =========================================================
|
||
|
|
# Codec
|
||
|
|
# =========================================================
|
||
|
|
|
||
|
|
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
|