475 lines
9.8 KiB
Python
475 lines
9.8 KiB
Python
import dataclasses
|
|
import json
|
|
import struct
|
|
from enum import auto, Enum
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
# from flandre.utils.RfMat import RfMat
|
|
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
|
|
|
|
|
|
class BG(Enum):
|
|
Msg1 = auto()
|
|
Msg2 = auto()
|
|
RefreshDeviceMsg = auto()
|
|
BMMsg = auto()
|
|
TickMsg = auto()
|
|
KillMsg = auto()
|
|
StrMsg = auto()
|
|
MoveAxisMsg = auto()
|
|
ImageArgMsg = auto()
|
|
SetSeqMetaMsg = auto()
|
|
SeqMetaMsg = auto()
|
|
SeqIdMinMax = auto()
|
|
SetBaseMsg = auto()
|
|
PlaybackSeqListMsg = auto()
|
|
SetPlayMode = auto()
|
|
SetDeviceEnabledMsg = auto()
|
|
SetDeviceConnectedMsg = auto()
|
|
DeviceEnabledMsg = auto()
|
|
DeviceConnectedMsg = auto()
|
|
SetDeviceConfigMsg = auto()
|
|
DeviceOnlineMsg = auto()
|
|
DeviceConfigListMsg = auto()
|
|
SetRecordMsg = auto()
|
|
RobotRtsiMsg = auto()
|
|
RecordFrameMsg = auto()
|
|
SeqIdList = auto()
|
|
SetWindowVisibleMsg = auto()
|
|
SetSidMsg = auto()
|
|
ImagingConfigNameListMsg = auto()
|
|
RfFrameMsg = auto()
|
|
RfFrameWithMetaMsg = auto()
|
|
BytesMsg = auto()
|
|
BeamformerMsg = auto()
|
|
RequestRfFrameMsg = auto()
|
|
MidiMsg = auto()
|
|
JoystickMsg = auto()
|
|
DeviceZero = auto()
|
|
SetDeviceSwitchMsg = auto()
|
|
DeviceSwitchMsg = auto()
|
|
HeaderByteMsg = auto()
|
|
RfMatMsg = auto()
|
|
KeyPressMsg = auto()
|
|
RGB888Msg = auto()
|
|
|
|
|
|
class Msg:
|
|
MAGIC = 1919810
|
|
|
|
@classmethod
|
|
def decode_base(cls, data: bytes) -> 'Msg':
|
|
return cls(**json.loads(data.decode()))
|
|
|
|
def encode(self) -> bytes:
|
|
return json.dumps(self.__dict__).encode()
|
|
|
|
@classmethod
|
|
def decode(cls, data: bytes) -> 'Msg':
|
|
return cls.decode_base(data)
|
|
|
|
@classmethod
|
|
def decode_msg(cls, msg: bytes):
|
|
magic, eid = struct.unpack('II', msg[:8])
|
|
assert magic == cls.MAGIC
|
|
class_: 'Msg' = globals()[BG(eid).name]
|
|
return class_.decode(msg[8:])
|
|
|
|
@classmethod
|
|
def eid(cls):
|
|
return struct.pack('I', BG[cls.__name__].value)
|
|
|
|
@classmethod
|
|
def magic(cls):
|
|
return struct.pack('I', cls.MAGIC)
|
|
|
|
def encode_msg(self):
|
|
return self.magic() + self.eid() + self.encode()
|
|
|
|
|
|
class HeaderByteMsg(Msg):
|
|
header: dict
|
|
data: bytes
|
|
|
|
def __init__(self, header: dict, data: bytes):
|
|
self.header = header
|
|
self.data = data
|
|
|
|
def encode(self) -> bytes:
|
|
e = json.dumps(self.header).encode()
|
|
return struct.pack('I', e.__len__()) + e + self.data
|
|
|
|
@classmethod
|
|
def decode(cls, data: bytes) -> 'HeaderByteMsg':
|
|
header_len = struct.unpack('I', data[:4])[0]
|
|
header = json.loads(data[4:4 + header_len])
|
|
data = data[4 + header_len:]
|
|
return HeaderByteMsg(header, data)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Msg1(Msg):
|
|
a: int = 0
|
|
b: int = 1
|
|
|
|
|
|
class RefreshDeviceMsg(Msg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class KillMsg(Msg):
|
|
name: str = ''
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Msg2(Msg):
|
|
a: int = 2
|
|
b: int = 2
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TickMsg(Msg):
|
|
time: float = 0
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class StrMsg(Msg):
|
|
value: str = ''
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class KeyPressMsg(StrMsg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class BoolMsg(Msg):
|
|
value: bool
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SetWindowVisibleMsg(Msg):
|
|
sender: str
|
|
name: str
|
|
value: bool
|
|
|
|
|
|
class SetDeviceEnabledMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
class SetDeviceConnectedMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
class DeviceEnabledMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
class DeviceConnectedMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
class DeviceOnlineMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SetRecordMsg(Msg):
|
|
enable: bool
|
|
commit: str = ''
|
|
base: str = ''
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SetDeviceConfigMsg(Msg):
|
|
name: str
|
|
txt: str
|
|
|
|
|
|
class SetBaseMsg(StrMsg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class PlaybackSeqListMsg(Msg):
|
|
value: list[str]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ImagingConfigNameListMsg(Msg):
|
|
value: list[str]
|
|
|
|
|
|
class SetPlayMode(StrMsg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SetSeqMetaMsg(Msg):
|
|
target: str
|
|
name: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SeqMetaMsg(Msg):
|
|
target: str
|
|
name: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SeqIdMinMax(Msg):
|
|
min: int
|
|
max: int
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SeqIdList(Msg):
|
|
li: list[int]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MoveAxisMsg(Msg):
|
|
sender: str
|
|
axis: str
|
|
value: int
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ImageArgMsg(Msg):
|
|
sender: str
|
|
t_end: int
|
|
t_start: int = 0
|
|
v2: int = 5900
|
|
dct_center: int = 0
|
|
dct_bandwidth: int = 0
|
|
f_rows: int = 0
|
|
beta: int = 10
|
|
|
|
@staticmethod
|
|
def from_path(p: Path):
|
|
return ImageArgMsg(
|
|
sender='load',
|
|
**json.loads(p.read_text())
|
|
)
|
|
|
|
def json(self):
|
|
arg_d = dict()
|
|
for field in dataclasses.fields(ImageArgMsg):
|
|
match field.name:
|
|
case 'sender':
|
|
pass
|
|
case _:
|
|
arg_d[field.name] = self.__getattribute__(field.name)
|
|
return json.dumps(arg_d)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class DeviceConfigListMsg(Msg):
|
|
arr: list[tuple[str, str]]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RecordFrameMsg(Msg):
|
|
size: int
|
|
current_sid: int
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class SetSidMsg(Msg):
|
|
value: int
|
|
|
|
|
|
class BMMsg(Msg):
|
|
def __init__(self, t: int, data: bytes):
|
|
self.data = data
|
|
self.t = t
|
|
|
|
def encode(self) -> bytes:
|
|
return struct.pack('I', self.t) + self.data
|
|
|
|
@classmethod
|
|
def decode(cls, data: bytes) -> 'Msg':
|
|
return cls(
|
|
struct.unpack('I', data[:4])[0],
|
|
data[4:]
|
|
)
|
|
|
|
|
|
class RfFrameMsg(Msg):
|
|
def __init__(self, sequence_id: int, encoder: int, data: bytes):
|
|
self.sequence_id = sequence_id
|
|
self.encoder = encoder
|
|
self.data = data
|
|
|
|
def encode(self) -> bytes:
|
|
return struct.pack(
|
|
'II',
|
|
self.sequence_id,
|
|
self.encoder,
|
|
) + self.data
|
|
|
|
@classmethod
|
|
def decode(cls, data: bytes) -> 'RfFrameMsg':
|
|
return cls(
|
|
*struct.unpack('II', data[:8]),
|
|
data[8:]
|
|
)
|
|
|
|
|
|
class RGB888Msg(HeaderByteMsg):
|
|
|
|
def __init__(self, b: bytes, w: int, h: int):
|
|
self.w = w
|
|
self.h = h
|
|
super().__init__(dict(h=h, w=w), b)
|
|
|
|
@classmethod
|
|
def decode(cls, data) -> 'RGB888Msg':
|
|
msg = super(RGB888Msg, cls).decode(data)
|
|
return RGB888Msg(msg.data, msg.header['w'], msg.header['h'])
|
|
|
|
|
|
class RfMatMsg(HeaderByteMsg):
|
|
|
|
def __init__(self, rfmat: 'RfMat'):
|
|
self.rfmat = rfmat
|
|
|
|
super().__init__(dict(
|
|
frame_meta=rfmat.frame_meta.name,
|
|
seq_meta=rfmat.seq_meta.name,
|
|
data_shape=rfmat.m.shape,
|
|
dtype=str(rfmat.m.dtype),
|
|
), rfmat.m.tobytes())
|
|
|
|
@classmethod
|
|
def decode(cls, data) -> 'RfMatMsg':
|
|
from flandre.utils.RfMat import RfMat
|
|
msg = super(RfMatMsg, cls).decode(data)
|
|
dt = np.dtype(msg.header['dtype'])
|
|
mat = np.frombuffer(msg.data, dtype=dt).reshape(msg.header['data_shape'])
|
|
rfmat = RfMat(
|
|
mat,
|
|
RfFrameMeta.from_name(msg.header['frame_meta']),
|
|
RfSequenceMeta.from_name(msg.header['seq_meta']),
|
|
)
|
|
return RfMatMsg(rfmat)
|
|
|
|
|
|
class RfFrameWithMetaMsg(HeaderByteMsg):
|
|
@dataclasses.dataclass
|
|
class Header:
|
|
sender: int
|
|
meta: str
|
|
is_zip: bool
|
|
|
|
def __init__(self, sender: int, meta: RfFrameMeta, data: bytes, is_zip: bool = False):
|
|
self._header = self.Header(sender, meta.name, is_zip)
|
|
super().__init__(self._header.__dict__, data)
|
|
|
|
@property
|
|
def sender(self) -> int:
|
|
return self._header.sender
|
|
|
|
@property
|
|
def meta(self) -> RfFrameMeta:
|
|
return RfFrameMeta.from_name(self._header.meta)
|
|
|
|
@classmethod
|
|
def decode(cls, data) -> 'RfFrameWithMetaMsg':
|
|
msg = super(RfFrameWithMetaMsg, cls).decode(data)
|
|
header = cls.Header(**msg.header)
|
|
return RfFrameWithMetaMsg(header.sender, RfFrameMeta.from_name(header.meta), msg.data, header.is_zip)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class BytesMsg(Msg):
|
|
value: bytes
|
|
|
|
def encode(self) -> bytes:
|
|
return self.value
|
|
|
|
@classmethod
|
|
def decode(cls, data: bytes) -> 'Msg':
|
|
return cls(data)
|
|
|
|
def split(self):
|
|
pass
|
|
# id2 = r.index(Msg.magic(), 1)
|
|
# id3 = r.index(Msg.magic(), id2 + 1)
|
|
# print(id2, id3)
|
|
# seq_msg: SetSeqMetaMsg = Msg.decode_msg(r[0:id2])
|
|
# arg_msg: ImageArgMsg = Msg.decode_msg(r[id2:id3])
|
|
# b_msg: RfFrameMsg = Msg.decode_msg(r[id3:-1])
|
|
|
|
|
|
class BeamformerMsg(BytesMsg):
|
|
pass
|
|
|
|
|
|
class DeviceZero(Msg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MidiMsg(Msg):
|
|
type: str
|
|
channel: int = None
|
|
pitch: int = None
|
|
control: int = None
|
|
value: int = None
|
|
velocity: int = None
|
|
note: int = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class JoystickMsg(Msg):
|
|
key: str
|
|
value: float
|
|
|
|
|
|
class DeviceSwitchMsg(StrMsg):
|
|
pass
|
|
|
|
|
|
class SetDeviceSwitchMsg(BoolMsg):
|
|
pass
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RobotRtsiMsg(Msg):
|
|
pos: tuple[int, int, int, int, int, int]
|
|
force: tuple[int, int, int, int, int, int]
|
|
|
|
|
|
class RequestRfFrameMsg(Msg):
|
|
pass
|
|
|
|
|
|
def test():
|
|
values = set(item.name for item in BG)
|
|
for k in globals().keys():
|
|
if k.endswith('Msg') and k not in ['Msg', 'BoolMsg']:
|
|
if k not in values:
|
|
raise RuntimeError(f"Unknown msg type: {k}")
|
|
|
|
|
|
test()
|
|
|
|
if __name__ == '__main__':
|
|
# c = HeaderByteMsg(dict(a=1, b='s'), b'asdasd')
|
|
# c2 = c.decode(c.encode())
|
|
# print(c2.header, c2.data)
|
|
|
|
# c = RfFrameWithMetaMsg(1, RfFrameMeta(1, 11, 111), b'asdasdasdads', True)
|
|
# print(c.sender, c.meta, c.data, c._header)
|
|
# c2 = c.decode(c.encode())
|
|
# print(c2.sender, c2.meta, c2.data, c2._header)
|
|
|
|
rfmat = RfMat(np.array([1, 2, 3]), RfFrameMeta(), RfSequenceMeta())
|
|
e = RfMatMsg(rfmat).encode()
|
|
d = RfMatMsg.decode(e)
|