agrobot_base/Python/raspi/multispectral_service.py

222 lines
6.7 KiB
Python

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"})