yukari/flandre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rm RfFrameWithMetaMsg

Browse Source 
add launch count
This commit is contained in:
flandre 2025-04-16 17:53:18 +08:00
parent 9a4d7eca5d
commit aa9141dc1a
15 changed files with 455 additions and 272 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
flandre/launcher.py
View File

@ -1,6 +1,7 @@
import logging import logging
import multiprocessing import multiprocessing
import os import os
import time
import tomllib import tomllib
from enum import Enum from enum import Enum
from pathlib import Path from pathlib import Path
@ -19,7 +20,7 @@ from flandre.nodes.ImageFFMPEG import ImageFFMPEG
from flandre.nodes.ImageQt import ImageQt from flandre.nodes.ImageQt import ImageQt
from flandre.nodes.Midi import Midi from flandre.nodes.Midi import Midi
from flandre.nodes.Mi import Mi from flandre.nodes.Mi import Mi
from flandre.utils.Msg import KillMsg from flandre.utils.Msg import KillMsg, TickMsg, Msg1
from flandre.config import CONFIG_FOLDER from flandre.config import CONFIG_FOLDER
@ -38,31 +39,40 @@ class LaunchComponent(Enum):
def launch(arg: dict[LaunchComponent, dict]): def launch(arg: dict[LaunchComponent, dict]):
ps = [ logging.basicConfig(level=logging.INFO)
Broker() multiprocessing.set_start_method('spawn')
] bp = multiprocessing.Process(target=Broker(broker=True))
bp.start()
# time.sleep(1)
ps = []
for k, v in arg.items(): for k, v in arg.items():
if k == LaunchComponent.MainUI and os.environ.get('XDG_CURRENT_DESKTOP', None) == 'KDE': if k == LaunchComponent.MainUI and os.environ.get('XDG_CURRENT_DESKTOP', None) == 'KDE':
ps.append(kde_pyqt6_mainui) ps.append(kde_pyqt6_mainui)
continue continue
ps.append(k.value(**v)) ps.append(k.value(**v))
logging.basicConfig(level=logging.INFO)
multiprocessing.set_start_method('spawn')
multiprocessing.Pool()
pps = [] pps = []
for p in ps: for p in ps:
pps.append(multiprocessing.Process(target=p)) pps.append(multiprocessing.Process(target=p))
for p in pps: for p in pps:
p.start() p.start()
c = BusClient(KillMsg) c = BusClient(KillMsg, TickMsg)
cnt = 0
while True: while True:
x: KillMsg = c.recv() msg = c.recv()
if x.name == '': if isinstance(msg, KillMsg):
if msg.name == '':
break break
if isinstance(msg, TickMsg):
cnt += 1
logging.error(f'{cnt}')
if cnt == len(ps):
c.send(Msg1())
for p in pps: for p in pps:
p.kill() p.kill()
bp.kill()
def launch_from_file(file: Path): def launch_from_file(file: Path):

20
flandre/nodes/Beamformer.py
View File

@ -10,10 +10,10 @@ from flandre.beamformer.dist import direct_dist
from flandre.config import C from flandre.config import C
from flandre.nodes.Node import Node from flandre.nodes.Node import Node
from flandre.utils.Config import DeviceConfig from flandre.utils.Config import DeviceConfig
from flandre.utils.Msg import ImageArgMsg, Msg, RfFrameWithMetaMsg, BeamformerMsg, \ from flandre.utils.Msg import ImageArgMsg, Msg, BeamformerMsg, \
SeqMetaMsg, RfMatMsg, RGB888Msg SeqMetaMsg, RfMatMsg, RGB888Msg, RfFrameMsg
from flandre.utils.RfFile import RfSequenceMeta
from flandre.utils.RfMat import RfMat from flandre.utils.RfMat import RfMat
from flandre.utils.RfMeta import RfSequenceMeta
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -54,7 +54,6 @@ class Beamformer(Node):
# self.send(RGB888Msg(d2.__bytes__(), d2.w, d2.h)) # self.send(RGB888Msg(d2.__bytes__(), d2.w, d2.h))
def loop(self): def loop(self):
time.sleep(1)
dc = DeviceConfig() dc = DeviceConfig()
pwi, _ = gen_pwi(direct_dist(dc), dc) pwi, _ = gen_pwi(direct_dist(dc), dc)
last_v2 = 5900 last_v2 = 5900
@ -73,21 +72,16 @@ class Beamformer(Node):
continue continue
r = msg.value r = msg.value
id2 = r.index(Msg.magic(), 1) id2 = r.index(Msg.magic(), 1)
id3 = r.index(Msg.magic(), id2 + 1) arg_msg: ImageArgMsg = Msg.decode_msg(r[:id2])
seq_msg: SeqMetaMsg = Msg.decode_msg(r[:id2]) rf_frame_msg: RfFrameMsg = Msg.decode_msg(r[id2:])
arg_msg: ImageArgMsg = Msg.decode_msg(r[id2:id3]) mat = RfMat.from_rf_frame(rf_frame_msg.rf_frame, 'gpu')
b_msg: RfFrameWithMetaMsg = Msg.decode_msg(r[id3:])
s = b_msg.data
fb1 = cp.frombuffer(s, dtype=cp.int16)
seq_meta = RfSequenceMeta.from_name(seq_msg.name)
mat = RfMat(fb1.reshape(seq_meta.shape), b_msg.meta, seq_meta)
if arg_msg.v2 != last_v2 or arg_msg.f_rows != last_f_rows: if arg_msg.v2 != last_v2 or arg_msg.f_rows != last_f_rows:
last_v2 = arg_msg.v2 last_v2 = arg_msg.v2
last_f_rows = arg_msg.f_rows last_f_rows = arg_msg.f_rows
dc = DeviceConfig(v2=arg_msg.v2, rows=arg_msg.f_rows) dc = DeviceConfig(v2=arg_msg.v2, rows=arg_msg.f_rows)
pwi, _ = gen_pwi(direct_dist(dc), dc) pwi, _ = gen_pwi(direct_dist(dc), dc)
try: try:
if seq_meta.mode == RfSequenceMeta.RfSequenceMode.PWI: if mat.seq_meta.mode == RfSequenceMeta.RfSequenceMode.PWI:
self.process_pwi(mat, arg_msg, pwi) self.process_pwi(mat, arg_msg, pwi)
except Exception as e: except Exception as e:
logger.warning(e) logger.warning(e)

10
flandre/nodes/Broker.py
View File

@ -4,17 +4,21 @@ import zmq
from zmq import ContextTerminated from zmq import ContextTerminated
from flandre.nodes.Node import Node from flandre.nodes.Node import Node
from flandre.utils.Msg import KillMsg from flandre.utils.Msg import KillMsg, TickMsg
class Broker(Node): class Broker(Node):
def loop(self): def loop(self):
def t(): def t():
while True: while True:
r:KillMsg = self.recv() msg = self.recv()
if r.name == '': if isinstance(msg, TickMsg):
pass
elif isinstance(msg, KillMsg):
if msg.name == '':
self.context.term() self.context.term()
break break
threading.Thread(target=t, daemon=True).start() threading.Thread(target=t, daemon=True).start()
frontend = self.context.socket(zmq.XSUB) frontend = self.context.socket(zmq.XSUB)
backend = self.context.socket(zmq.XPUB) backend = self.context.socket(zmq.XPUB)

26
flandre/nodes/Device.py
View File

@ -11,9 +11,10 @@ from flandre.BusClient import BusClient
from flandre.config import C from flandre.config import C
from flandre.nodes.Node import Node from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, KillMsg, SetDeviceConnectedMsg, SetDeviceEnabledMsg, DeviceEnabledMsg, \ from flandre.utils.Msg import ImageArgMsg, KillMsg, SetDeviceConnectedMsg, SetDeviceEnabledMsg, DeviceEnabledMsg, \
DeviceConnectedMsg, SetDeviceConfigMsg, DeviceOnlineMsg, DeviceConfigListMsg, RequestRfFrameMsg, RfFrameWithMetaMsg, \ DeviceConnectedMsg, SetDeviceConfigMsg, DeviceOnlineMsg, DeviceConfigListMsg, RequestRfFrameMsg, \
DeviceZero, SetDeviceSwitchMsg, DeviceSwitchMsg, SeqMetaMsg, RefreshDeviceMsg DeviceZero, SetDeviceSwitchMsg, DeviceSwitchMsg, SeqMetaMsg, RefreshDeviceMsg, RfFrameMsg
from flandre.utils.RfMeta import RfFrameMeta from flandre.utils.RfFrame import RfFrameMemory
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -43,7 +44,7 @@ class Device(Node):
def __init__(self, level=logging.INFO): def __init__(self, level=logging.INFO):
super(Device, self).__init__(level=level) super(Device, self).__init__(level=level)
self.arg = ImageArgMsg('', t_start=0, t_end=1499) self.arg = ImageArgMsg('', t_start=0, t_end=1499)
self.seq_meta = None self.seq_meta: RfSequenceMeta | None = None
self.req_driver_socket: zmq.Socket = None self.req_driver_socket: zmq.Socket = None
self.rep_socket: zmq.Socket = None self.rep_socket: zmq.Socket = None
self.ok = b'ok\x00' self.ok = b'ok\x00'
@ -147,8 +148,10 @@ class Device(Node):
def get_seq_meta_name(self): def get_seq_meta_name(self):
self.device_cmd(DeviceCmd.GetName) self.device_cmd(DeviceCmd.GetName)
rb = self.req_driver_socket.recv() rb = self.req_driver_socket.recv()
name = rb.decode()
self.seq_meta = RfSequenceMeta.from_name(name)
if rb != b'': if rb != b'':
self.send(SeqMetaMsg('live', rb.decode())) self.send(SeqMetaMsg('live', name))
def get_connection(self): def get_connection(self):
self.device_cmd(DeviceCmd.GetConnection) self.device_cmd(DeviceCmd.GetConnection)
@ -235,6 +238,7 @@ class Device(Node):
self.disconnect() self.disconnect()
elif isinstance(msg, SetDeviceConfigMsg): elif isinstance(msg, SetDeviceConfigMsg):
self.set_name_and_file(msg.name, msg.txt) self.set_name_and_file(msg.name, msg.txt)
self.seq_meta = RfSequenceMeta.from_name(msg.name)
self.send(SeqMetaMsg('live', msg.name)) self.send(SeqMetaMsg('live', msg.name))
elif isinstance(msg, RequestRfFrameMsg): elif isinstance(msg, RequestRfFrameMsg):
if self.switch: if self.switch:
@ -244,10 +248,14 @@ class Device(Node):
continue continue
_, sequence_id, encoder = struct.unpack_from('=iqi', braw) _, sequence_id, encoder = struct.unpack_from('=iqi', braw)
buffer = braw[4 + 8 + 4:] buffer = braw[4 + 8 + 4:]
# logger.debug('send') if self.seq_meta is not None:
self.send(RfFrameWithMetaMsg(0, RfFrameMeta( self.send(RfFrameMsg(0, RfFrameMemory(
encoder=encoder, sequence_id=sequence_id RfFrameMeta(encoder=encoder, sequence_id=sequence_id),
), buffer)) self.seq_meta,
buffer,
)))
else:
logger.warning(f'no seq meta msg')
else: else:
logger.warning('device not online') logger.warning('device not online')
elif isinstance(msg, DeviceZero): elif isinstance(msg, DeviceZero):

12
flandre/nodes/Loader.py
View File

@ -7,8 +7,8 @@ import zmq
from flandre.config import C, ISDEV from flandre.config import C, ISDEV
from flandre.nodes.Node import Node from flandre.nodes.Node import Node
from flandre.utils.Msg import MoveAxisMsg, KillMsg, SetSeqMetaMsg, SeqIdMinMax, SetBaseMsg, PlaybackSeqListMsg, \ from flandre.utils.Msg import MoveAxisMsg, KillMsg, SetSeqMetaMsg, SeqIdMinMax, SetBaseMsg, PlaybackSeqListMsg, \
SeqIdList, SetSidMsg, RfFrameWithMetaMsg SeqIdList, SetSidMsg, RfFrameMsg
from flandre.utils.RfFile import RfSequence from flandre.utils.RfSequence import RfSequence
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -30,18 +30,14 @@ class Loader(Node):
if msg.axis == 'S': if msg.axis == 'S':
pass pass
elif isinstance(msg, SetSidMsg): elif isinstance(msg, SetSidMsg):
frame = rff.frames[msg.value] self.send(RfFrameMsg(1, rff.frames[msg.value]))
self.send(RfFrameWithMetaMsg(1, frame.meta, frame.bytes, frame.is_zip))
elif isinstance(msg, SetSeqMetaMsg): elif isinstance(msg, SetSeqMetaMsg):
if base is None: if base is None:
continue continue
if msg.target == 'playback': if msg.target == 'playback':
logger.info(f'load {msg.name}') logger.info(f'load {msg.name}')
if msg.name.endswith('.zip'): rff = RfSequence(base / msg.name)
rff = RfSequence.from_zip(base / msg.name)
else:
rff = RfSequence.from_folder(base / msg.name)
self.send(SeqIdMinMax(*rff.seq_id_minmax)) self.send(SeqIdMinMax(*rff.seq_id_minmax))
self.send(SeqIdList([f.meta.sequence_id for f in rff.frames])) self.send(SeqIdList([f.meta.sequence_id for f in rff.frames]))
self.send(SetSidMsg(0)) self.send(SetSidMsg(0))

93
flandre/nodes/Muxer.py
View File

@ -8,27 +8,26 @@ import zmq
from flandre.config import C from flandre.config import C
from flandre.nodes.Node import Node from flandre.nodes.Node import Node
from flandre.utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, \ from flandre.utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, \
ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg ImagingConfigNameListMsg, BeamformerMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg, RfFrameMsg
from flandre.utils.RfFile import RfSequenceMeta from flandre.utils.RfFrame import RfFrameMemory
from flandre.utils.RfMeta import RfFrameMeta from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
from flandre.utils.archive import zip_to_bytes
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class Muxer(Node): class Muxer(Node):
topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, RfFrameWithMetaMsg, topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, RfFrameMsg,
ImageArgMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg] ImageArgMsg, RobotRtsiMsg, SeqMetaMsg, DeviceEnabledMsg]
def __init__(self, level=logging.INFO): def __init__(self, level=logging.INFO):
super(Muxer, self).__init__(level=level) super(Muxer, self).__init__(level=level)
self.seq_meta = None self.seq_meta = None
self.seq_meta_live = None self.seq_meta_live: RfSequenceMeta = None
self.seq_meta_playback = None self.seq_meta_playback = None
self.play_mode = None self.play_mode = None
self.rep_socket: zmq.Socket = None self.rep_socket: zmq.Socket = None
self.req_driver_socket: zmq.Socket = None self.req_driver_socket: zmq.Socket = None
self.playback_rf_msg: RfFrameWithMetaMsg | None = None self.playback_rf_msg: RfFrameMsg | None = None
self.device_enabled = False self.device_enabled = False
def custom_setup(self): def custom_setup(self):
@ -38,14 +37,6 @@ class Muxer(Node):
self.req_driver_socket.connect(C.driver_rep_socket) self.req_driver_socket.connect(C.driver_rep_socket)
self.c.poller.register(self.rep_socket, zmq.POLLIN) self.c.poller.register(self.rep_socket, zmq.POLLIN)
@property
def current_meta(self):
seq_meta: RfSequenceMeta | None = {
"live": self.seq_meta_live,
"playback": self.seq_meta_playback,
}.get(self.play_mode, None)
return seq_meta
def handle_rep_socket(self): def handle_rep_socket(self):
self.rep_socket.recv() self.rep_socket.recv()
if self.play_mode is None: if self.play_mode is None:
@ -53,22 +44,11 @@ class Muxer(Node):
return return
match self.play_mode: match self.play_mode:
case 'playback': case 'playback':
# logger.warning(f'test, {self.playback_rf_msg}')
if self.playback_rf_msg is None: if self.playback_rf_msg is None:
self.rep_socket.send(BeamformerMsg(b'nop').encode_msg()) self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
return return
data_msg = self.playback_rf_msg data_msg = self.playback_rf_msg
# logger.info(f'bb {data_msg._header.is_zip}')
if data_msg._header.is_zip:
p = Path(data_msg.data.decode())
if p.suffix == '.zst':
data_msg = RfFrameWithMetaMsg(
data_msg.sender,
data_msg.meta,
zip_to_bytes(p.parent, int(p.stem))
)
else:
raise NotImplementedError()
case 'live': case 'live':
if not self.device_enabled: if not self.device_enabled:
self.rep_socket.send(BeamformerMsg(b'init').encode_msg()) self.rep_socket.send(BeamformerMsg(b'init').encode_msg())
@ -82,17 +62,32 @@ class Muxer(Node):
return return
_, _, sequence_id, encoder = struct.unpack_from('=iiii', driver_data_raw) _, _, sequence_id, encoder = struct.unpack_from('=iiii', driver_data_raw)
driver_data_body = driver_data_raw[4 + 8 + 4:] driver_data_body = driver_data_raw[4 + 8 + 4:]
# print(driver_data_body.__len__()/256/2) data_msg = RfFrameMsg(0, RfFrameMemory(
data_msg = RfFrameWithMetaMsg(0, RfFrameMeta( RfFrameMeta(
encoder=encoder, sequence_id=sequence_id encoder=encoder,
), driver_data_body) sequence_id=sequence_id,
robot_x=self.rtsi.pos[0],
robot_y=self.rtsi.pos[1],
robot_z=self.rtsi.pos[2],
robot_roll=self.rtsi.pos[3],
robot_pitch=self.rtsi.pos[4],
robot_yal=self.rtsi.pos[5],
robot_force_x=self.rtsi.force[0],
robot_force_y=self.rtsi.force[1],
robot_force_z=self.rtsi.force[2],
robot_force_roll=self.rtsi.force[3],
robot_force_pitch=self.rtsi.force[4],
robot_force_yal=self.rtsi.force[5],
),
self.seq_meta_live,
driver_data_body
))
case _: case _:
raise NotImplementedError() raise NotImplementedError()
if (data_msg.data.__len__() // 2) != self.current_meta.prod(): # if (data_msg.data.__len__() // 2) != data_msg.rf_frame.prod():
self.rep_socket.send(BeamformerMsg(b'nop').encode_msg()) # self.rep_socket.send(BeamformerMsg(b'nop').encode_msg())
return # return
self.rep_socket.send(BeamformerMsg( self.rep_socket.send(BeamformerMsg(
SeqMetaMsg('any', self.current_meta.name).encode_msg() +
self.arg.encode_msg() + self.arg.encode_msg() +
data_msg.encode_msg() data_msg.encode_msg()
).encode_msg()) ).encode_msg())
@ -106,7 +101,6 @@ class Muxer(Node):
self.arg = ImageArgMsg('', t_start=0, t_end=1499) self.arg = ImageArgMsg('', t_start=0, t_end=1499)
self.c.poller.register(device_socket, zmq.POLLIN) self.c.poller.register(device_socket, zmq.POLLIN)
time.sleep(1)
self.send(ImagingConfigNameListMsg([path.stem for path in C.imaging_config_folder.glob('*.json')])) self.send(ImagingConfigNameListMsg([path.stem for path in C.imaging_config_folder.glob('*.json')]))
while True: while True:
socks = dict(self.c.poller.poll()) socks = dict(self.c.poller.poll())
@ -117,31 +111,12 @@ class Muxer(Node):
self.handle_rep_socket() self.handle_rep_socket()
if k == self.c.sub: if k == self.c.sub:
msg = self.recv() msg = self.recv()
if isinstance(msg, RfFrameWithMetaMsg): if isinstance(msg, KillMsg):
if msg.sender == 0 and self.play_mode == 'live':
msg.meta.robot_x = self.rtsi.pos[0]
msg.meta.robot_y = self.rtsi.pos[1]
msg.meta.robot_z = self.rtsi.pos[2]
msg.meta.robot_roll = self.rtsi.pos[3]
msg.meta.robot_pitch = self.rtsi.pos[4]
msg.meta.robot_yal = self.rtsi.pos[5]
msg.meta.robot_force_x = self.rtsi.force[0]
msg.meta.robot_force_y = self.rtsi.force[1]
msg.meta.robot_force_z = self.rtsi.force[2]
msg.meta.robot_force_roll = self.rtsi.force[3]
msg.meta.robot_force_pitch = self.rtsi.force[4]
msg.meta.robot_force_yal = self.rtsi.force[5]
self.send(BeamformerMsg(
SetSeqMetaMsg('any', self.seq_meta_live.name).encode_msg() +
self.arg.encode_msg() +
msg.encode_msg()
))
elif msg.sender == 1:
self.playback_rf_msg = msg
elif isinstance(msg, KillMsg):
if msg.name == '': if msg.name == '':
return return
elif isinstance(msg, RfFrameMsg):
if msg.sender == 1:
self.playback_rf_msg = msg
elif isinstance(msg, ImageArgMsg): elif isinstance(msg, ImageArgMsg):
self.arg = msg self.arg = msg
elif isinstance(msg, SetSeqMetaMsg): elif isinstance(msg, SetSeqMetaMsg):

12
flandre/nodes/Node.py
View File

@ -1,10 +1,11 @@
import logging import logging
import time
from abc import abstractmethod from abc import abstractmethod
import zmq import zmq
from flandre.BusClient import BusClient from flandre.BusClient import BusClient
from flandre.utils.Msg import Msg, KillMsg from flandre.utils.Msg import Msg, KillMsg, TickMsg, Msg1
class Node: class Node:
@ -12,11 +13,12 @@ class Node:
bp = BusClient.bp bp = BusClient.bp
topics = [] topics = []
def __init__(self, enable_init=True, level=logging.INFO, conflare=False): def __init__(self, enable_init=True, level=logging.INFO, conflare=False, broker=False):
self.enable_init = enable_init self.enable_init = enable_init
self.isalive = True self.isalive = True
self.level = level self.level = level
self.conflare = conflare self.conflare = conflare
self.broker = broker
def recv(self): def recv(self):
return self.c.recv() return self.c.recv()
@ -53,9 +55,13 @@ class Node:
self.context = zmq.Context() self.context = zmq.Context()
if self.enable_init: if self.enable_init:
self.c = BusClient(*([KillMsg] + self.topics), poller=True, conflare=self.conflare) self.c = BusClient(*([KillMsg, Msg1] + self.topics), poller=True, conflare=self.conflare)
def __call__(self, *args, **kwargs): def __call__(self, *args, **kwargs):
self.setup() self.setup()
if not self.broker:
time.sleep(1)
self.send(TickMsg())
self.recv()
self.loop() self.loop()
print(self.__class__.__name__, 'exiting') print(self.__class__.__name__, 'exiting')

1
flandre/nodes/Recorder.py
View File

@ -35,7 +35,6 @@ class Recorder(Node):
) )
device_socket = self.context.socket(zmq.PULL) device_socket = self.context.socket(zmq.PULL)
self.c.poller.register(device_socket, zmq.POLLIN) self.c.poller.register(device_socket, zmq.POLLIN)
time.sleep(1)
while True: while True:
socks = dict(self.c.poller.poll()) socks = dict(self.c.poller.poll())
for k in socks: for k in socks:

114
flandre/utils/Msg.py
View File

@ -6,8 +6,9 @@ from pathlib import Path
import numpy as np import numpy as np
# from flandre.utils.RfMat import RfMat from flandre.utils.RfFrame import RfFrame, RfFrameMemory, RfFrameFile
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
from flandre.utils.RfSequence import RfSequence
class BG(Enum): class BG(Enum):
@ -40,8 +41,7 @@ class BG(Enum):
SetWindowVisibleMsg = auto() SetWindowVisibleMsg = auto()
SetSidMsg = auto() SetSidMsg = auto()
ImagingConfigNameListMsg = auto() ImagingConfigNameListMsg = auto()
RfFrameMsg = auto()
RfFrameWithMetaMsg = auto()
BytesMsg = auto() BytesMsg = auto()
BeamformerMsg = auto() BeamformerMsg = auto()
RequestRfFrameMsg = auto() RequestRfFrameMsg = auto()
@ -51,7 +51,11 @@ class BG(Enum):
SetDeviceSwitchMsg = auto() SetDeviceSwitchMsg = auto()
DeviceSwitchMsg = auto() DeviceSwitchMsg = auto()
HeaderByteMsg = auto() HeaderByteMsg = auto()
RfFrameMsg = auto()
# RfFrameWithMetaMsg = auto()
RfMatMsg = auto() RfMatMsg = auto()
KeyPressMsg = auto() KeyPressMsg = auto()
RGB888Msg = auto() RGB888Msg = auto()
@ -96,6 +100,8 @@ class HeaderByteMsg(Msg):
def __init__(self, header: dict, data: bytes): def __init__(self, header: dict, data: bytes):
self.header = header self.header = header
self.data = data self.data = data
if self.data is None:
self.data = b''
def encode(self) -> bytes: def encode(self) -> bytes:
e = json.dumps(self.header).encode() e = json.dumps(self.header).encode()
@ -299,25 +305,47 @@ class BMMsg(Msg):
) )
class RfFrameMsg(Msg): class RfFrameMsg(HeaderByteMsg):
def __init__(self, sequence_id: int, encoder: int, data: bytes): def __init__(self, sender: int, rf_frame: RfFrame):
self.sequence_id = sequence_id self.sender = sender
self.encoder = encoder self.rf_frame = rf_frame
self.data = data if isinstance(rf_frame, RfFrameFile):
super().__init__(dict(
def encode(self) -> bytes: sender=sender,
return struct.pack( type='RfFrameFile',
'II', frame_meta=rf_frame.meta.name,
self.sequence_id, frame_filename=rf_frame.filename,
self.encoder, seq_path=rf_frame.seq.path.__str__(),
) + self.data ), rf_frame.data)
elif isinstance(rf_frame, RfFrameMemory):
super().__init__(dict(
sender=sender,
type='RfFrameMemory',
frame_meta=rf_frame.meta.name,
seq_meta=rf_frame.seq_meta,
), rf_frame.data)
@classmethod @classmethod
def decode(cls, data: bytes) -> 'RfFrameMsg': def decode(cls, data: bytes) -> 'RfFrameMsg':
return cls( msg = super(RfFrameMsg, cls).decode(data)
*struct.unpack('II', data[:8]), if msg.header['type'] == 'RfFrameFile':
data[8:] return RfFrameMsg(
) msg.header['sender'],
RfFrameFile(
RfFrameMeta.from_name(msg.header['frame_meta']),
RfSequence(msg.header['seq_path']),
msg.data if msg.data.__len__() > 0 else None,
msg.header['frame_filename']
))
elif msg.header['type'] == 'RfFrameMemory':
return RfFrameMsg(
msg.header['sender'],
RfFrameMemory(
RfFrameMeta.from_name(msg.header['frame_meta']),
RfSequenceMeta.from_name(msg.header['seq_meta']),
msg.data,
))
raise NotImplementedError()
class RGB888Msg(HeaderByteMsg): class RGB888Msg(HeaderByteMsg):
@ -359,30 +387,30 @@ class RfMatMsg(HeaderByteMsg):
return RfMatMsg(rfmat) return RfMatMsg(rfmat)
class RfFrameWithMetaMsg(HeaderByteMsg): # class RfFrameWithMetaMsg(HeaderByteMsg):
@dataclasses.dataclass # @dataclasses.dataclass
class Header: # class Header:
sender: int # sender: int
meta: str # meta: str
is_zip: bool # is_zip: bool
#
def __init__(self, sender: int, meta: RfFrameMeta, data: bytes, is_zip: bool = False): # def __init__(self, sender: int, meta: RfFrameMeta, data: bytes, is_zip: bool = False):
self._header = self.Header(sender, meta.name, is_zip) # self._header = self.Header(sender, meta.name, is_zip)
super().__init__(self._header.__dict__, data) # super().__init__(self._header.__dict__, data)
#
@property # @property
def sender(self) -> int: # def sender(self) -> int:
return self._header.sender # return self._header.sender
#
@property # @property
def meta(self) -> RfFrameMeta: # def meta(self) -> RfFrameMeta:
return RfFrameMeta.from_name(self._header.meta) # return RfFrameMeta.from_name(self._header.meta)
#
@classmethod # @classmethod
def decode(cls, data) -> 'RfFrameWithMetaMsg': # def decode(cls, data) -> 'RfFrameWithMetaMsg':
msg = super(RfFrameWithMetaMsg, cls).decode(data) # msg = super(RfFrameWithMetaMsg, cls).decode(data)
header = cls.Header(**msg.header) # header = cls.Header(**msg.header)
return RfFrameWithMetaMsg(header.sender, RfFrameMeta.from_name(header.meta), msg.data, header.is_zip) # return RfFrameWithMetaMsg(header.sender, RfFrameMeta.from_name(header.meta), msg.data, header.is_zip)
@dataclasses.dataclass @dataclasses.dataclass

227
flandre/utils/RfFile.py
View File

@ -1,109 +1,118 @@
import json # import json
import zipfile # import zipfile
from pathlib import Path # from pathlib import Path
#
import cupy as cp # import cupy as cp
import numpy as np # import numpy as np
#
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta # from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
#
#
class RfFrame: # class RfFrame:
def __init__(self, data: bytes | Path, meta: RfFrameMeta, seq: 'RfSequence' = None): # def __init__(self, data: bytes | Path, meta: RfFrameMeta, seq: 'RfSequence' = None):
self.data = data # self.data = data
self.meta = meta # self.meta = meta
self._seq = seq # self._seq = seq
self.is_zip = False # self.is_zip = False
if isinstance(self.data, Path): # if isinstance(self.data, Path):
if self.data.parent.suffix == '.zip': # if self.data.parent.suffix == '.zip':
self.is_zip = True # self.is_zip = True
#
def save(self, folder: Path): # def save(self, folder: Path):
(folder / self.meta.name).write_bytes(self.bytes) # (folder / self.meta.name).write_bytes(self.bytes)
#
@property # @property
def seq(self): # def seq(self):
if self._seq is None and isinstance(self.data, Path): # if self._seq is None and isinstance(self.data, Path):
return RfSequence.from_folder(self.data.parent) # return RfSequence.from_folder(self.data.parent)
return self._seq # return self._seq
#
@property # @property
def bytes(self): # def bytes(self):
if self.is_zip: # if self.is_zip:
return self.data.__str__().encode() # return self.data.__str__().encode()
if isinstance(self.data, bytes): # if isinstance(self.data, bytes):
return self.data # return self.data
if isinstance(self.data, Path): # if isinstance(self.data, Path):
return self.data.read_bytes() # return self.data.read_bytes()
def mat(self, device='gpu'): #
from flandre.utils.RfMat import RfMat # def mat(self, device='gpu'):
# from flandre.utils.RfMat import RfMat
if device == 'gpu': #
arr = cp.frombuffer(self.bytes, dtype=cp.int16) # if device == 'gpu':
else: # arr = cp.frombuffer(self.bytes, dtype=cp.int16)
arr = np.frombuffer(self.bytes, dtype=np.int16) # else:
return RfMat(arr.reshape(self.seq.meta.shape), frame_meta=self.meta, seq_meta=self.seq.meta) # arr = np.frombuffer(self.bytes, dtype=np.int16)
# return RfMat(arr.reshape(self.seq.meta.shape), frame_meta=self.meta, seq_meta=self.seq.meta)
#
class RfSequence: #
def __init__(self, frames: list[RfFrame], meta: RfSequenceMeta): # class RfSequence:
self.frames = frames # def __init__(self, frames: list[RfFrame], meta: RfSequenceMeta):
self.meta = meta # self.frames = frames
# self.meta = meta
@classmethod #
def from_folder(cls, folder: Path | str) -> 'RfSequence': # @classmethod
folder = Path(folder) # def from_folder(cls, folder: Path | str) -> 'RfSequence':
if not folder.exists(): # folder = Path(folder)
raise FileNotFoundError # if not folder.exists():
meta = RfSequenceMeta.from_path(folder) # raise FileNotFoundError
rs = RfSequence([], meta) # meta = RfSequenceMeta.from_path(folder)
for f in folder.glob('*.bin'): # rs = RfSequence([], meta)
rs.frames.append(RfFrame(f, RfFrameMeta.from_path(f), seq=rs)) # for f in folder.glob('*.bin'):
return rs # rs.frames.append(RfFrame(f, RfFrameMeta.from_path(f), seq=rs))
# return rs
@classmethod #
def from_zip(cls, z: Path | str) -> 'RfSequence': # @classmethod
z = Path(z) # def from_zip(cls, z: Path | str) -> 'RfSequence':
if not z.exists(): # z = Path(z)
raise FileNotFoundError # if not z.exists():
meta = RfSequenceMeta.from_name(z.stem) # raise FileNotFoundError
rs = RfSequence([], meta) # meta = RfSequenceMeta.from_name(z.stem)
zz = zipfile.ZipFile(z) # rs = RfSequence([], meta)
max_idx = max([int(Path(item.filename).stem) for item in zz.infolist()]) # zz = zipfile.ZipFile(z)
for i in range(max_idx + 1): # max_idx = max([int(Path(item.filename).stem) for item in zz.infolist()])
j = json.loads(zz.open(f'{i}.json').read()) # for i in range(max_idx + 1):
rs.frames.append(RfFrame(z / f'{i}.zst', RfFrameMeta(**j), seq=rs)) # j = json.loads(zz.open(f'{i}.json').read())
return rs # rs.frames.append(RfFrame(z / f'{i}.zst', RfFrameMeta(**j), seq=rs))
# return rs
@property #
def all(self): # @property
pass # def all(self):
# pass
def query(self): #
pass # def query(self):
# pass
@property #
def seq_id_minmax(self): # @property
mmin = 2 ** 32 # def seq_id_minmax(self):
mmax = 0 # mmin = 2 ** 32
for f in self.frames: # mmax = 0
mmin = min(mmin, f.meta.sequence_id) # for f in self.frames:
mmax = max(mmax, f.meta.sequence_id) # mmin = min(mmin, f.meta.sequence_id)
return mmin, mmax # mmax = max(mmax, f.meta.sequence_id)
# return mmin, mmax
#
if __name__ == '__main__': #
# t = (1, 2) #
# f = RfSequenceMeta.from_name('123123,U=321,S=(1 2 3),M=PWI') #
# print(f.commit) #
# print(f.name) # if __name__ == '__main__':
# print(RfSequence.RfSequenceMeta.p2t) # f = RfSequence2('/mnt/16T/private_dataset/us/steel-top,U=30,M=PWI,S=(256 1502).zip').frames
# f = RfFrame.RfFrameMeta(123, 345) # for f2 in f:
# print(f.name) # print(f2.__bytes__().__len__())
# rs = RfSequence([], RfSequenceMeta()) # # print(Path('xxx/asdasdasd_asd=a1,ds1.123').with_suffix(''))
rs = RfSequence.from_folder( # # print(Path('xxx/asdasdasd_asd=a1,ds1.123').stem)
'/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/cap4/trim/R1.1,U=30,M=PWI,S=(256 1502)') # # t = (1, 2)
# print(rs.meta) # # f = RfSequenceMeta.from_name('123123,U=321,S=(1 2 3),M=PWI')
for frame in rs.frames: # # print(f.commit)
if frame.mat().rotate90().show((1080, 1920)) == ord('q'): # # print(f.name)
break # # print(RfSequence.RfSequenceMeta.p2t)
# # f = RfFrame.RfFrameMeta(123, 345)
# # print(f.name)
# # rs = RfSequence([], RfSequenceMeta())
# # rs = RfSequence.from_folder(
# # '/run/media/lambda/b86dccdc-f134-464b-a310-6575ee9ae85c/cap4/trim/R1.1,U=30,M=PWI,S=(256 1502)')
# # # print(rs.meta)
# # for frame in rs.frames:
# # if frame.mat().rotate90().show((1080, 1920)) == ord('q'):
# # break

39
flandre/utils/RfFrame.py Normal file
View File

@ -0,0 +1,39 @@
from pathlib import Path
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
from flandre.utils.RfSequence import RfSequence
class RfFrame:
def __bytes__(self) -> bytes:
raise NotImplementedError()
class RfFrameFile(RfFrame):
def __init__(self, meta: RfFrameMeta, seq: RfSequence, data: bytes = None, filename: str = None):
self.seq = seq
self.meta = meta
self.data = data
self.filename = filename
def __bytes__(self):
if self.data is not None:
return self.data
match self.seq.type:
case 'zip':
from flandre.utils.archive import zip_to_bytes
if self.filename is not None:
return zip_to_bytes(self.seq.path, int(Path(self.filename).stem))
case 'dir':
return (self.seq.path / self.filename).read_bytes()
raise NotImplementedError()
class RfFrameMemory(RfFrame):
def __init__(self, meta: RfFrameMeta, seq_meta: RfSequenceMeta, data: bytes):
self.seq_meta = seq_meta
self.meta = meta
self.data = data
def __bytes__(self):
return self.data

19
flandre/utils/RfMat.py
View File

@ -16,8 +16,8 @@ import scipy.signal
from cupyx.scipy.fft import dctn, idctn from cupyx.scipy.fft import dctn, idctn
from scipy.stats import norm as norms from scipy.stats import norm as norms
from flandre.utils.RfFile import RfFrame, RfSequenceMeta from flandre.utils.RfFrame import RfFrame, RfFrameFile, RfFrameMemory
from flandre.utils.RfMeta import RfFrameMeta from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
def hsv_to_rgb(hsv): def hsv_to_rgb(hsv):
@ -123,6 +123,19 @@ def bypassClass(original_class):
@bypassClass @bypassClass
class RfMat: class RfMat:
@staticmethod
def from_rf_frame(frame: RfFrame, device='cpu'):
if isinstance(frame, RfFrameFile):
seq_meta = frame.seq.meta
elif isinstance(frame, RfFrameMemory):
seq_meta = frame.seq_meta
else:
raise NotImplementedError()
m = np.frombuffer(frame.__bytes__(), dtype=np.int16).reshape(seq_meta.shape)
if device == 'gpu':
m = cp.asarray(m)
return RfMat(m, frame.meta, seq_meta)
def __init__(self, def __init__(self,
data: cp.ndarray, data: cp.ndarray,
frame_meta: RfFrameMeta = None, frame_meta: RfFrameMeta = None,
@ -339,6 +352,7 @@ class RfMat:
addend = self.p.zeros((1, h), dtype=self.p.float32) addend = self.p.zeros((1, h), dtype=self.p.float32)
addend[:, start:] = (self.p.arange(h - start) * scale) + 1 addend[:, start:] = (self.p.arange(h - start) * scale) + 1
return self.copy(self.m * addend) return self.copy(self.m * addend)
def time_gain_compensation_linear(self, scale: float, start: int = 0): def time_gain_compensation_linear(self, scale: float, start: int = 0):
h = self.m.shape[-1] h = self.m.shape[-1]
addend = self.p.zeros((1, h), dtype=np.int64) addend = self.p.zeros((1, h), dtype=np.int64)
@ -363,5 +377,6 @@ class RfMat:
self.m = self.m.astype(np.int64) self.m = self.m.astype(np.int64)
return self return self
if __name__ == '__main__': if __name__ == '__main__':
cp.zeros((1, 2, 3)) + 1 cp.zeros((1, 2, 3)) + 1

47
flandre/utils/RfSequence.py Normal file
View File

@ -0,0 +1,47 @@
import zipfile
from pathlib import Path
from flandre.utils.RfMeta import RfFrameMeta, RfSequenceMeta
class RfSequence:
def __init__(self, path: Path | str):
from flandre.utils.RfFrame import RfFrameFile
path = Path(path)
if not path.exists():
raise FileNotFoundError()
if path.is_dir():
self.type = 'dir'
elif path.suffix == '.zip':
self.type = 'zip'
self.path = path
self.meta = RfSequenceMeta.from_name(path.stem)
self._frames: list[RfFrameFile] | None = None
@property
def frames(self):
from flandre.utils.RfFrame import RfFrameFile
if self._frames is None:
arr = []
match self.type:
case 'zip':
zip_file = zipfile.ZipFile(self.path)
max_idx = max([int(Path(item.filename).stem) for item in zip_file.infolist()])
for i in range(max_idx + 1):
meta = RfFrameMeta.from_name(zip_file.open(f'{i}.meta').read().decode())
arr.append(RfFrameFile(meta, seq=self, filename=f'{i}.zst'))
case 'dir':
for f in self.path.glob('*.bin'):
meta = RfFrameMeta.from_name(f.stem)
arr.append(RfFrameFile(meta, seq=self, filename=f.name))
self._frames = arr
return self._frames
@property
def seq_id_minmax(self):
mmin = 2 ** 32
mmax = 0
for f in self.frames:
mmin = min(mmin, f.meta.sequence_id)
mmax = max(mmax, f.meta.sequence_id)
return mmin, mmax

9
flandre/utils/archive.py
View File

@ -42,8 +42,8 @@ def to_zip(li: list[tuple[Path, 'RfFrameMeta', list[tuple[Path, str]]]], temp_ds
dst = temp_dst / f'{i}.zst' dst = temp_dst / f'{i}.zst'
subprocess.run(['zstd', '-f', src, '-o', dst]) subprocess.run(['zstd', '-f', src, '-o', dst])
dstj = temp_dst / f'{i}.json' dstj = temp_dst / f'{i}.meta'
dstj.write_text(meta.json_str) dstj.write_text(meta.name)
for srcf, suffix in farr: for srcf, suffix in farr:
dstf = temp_dst / f'{i}{suffix}' dstf = temp_dst / f'{i}{suffix}'
@ -54,9 +54,12 @@ def to_zip(li: list[tuple[Path, 'RfFrameMeta', list[tuple[Path, str]]]], temp_ds
def zip_to_bytes(file: Path, name: int): def zip_to_bytes(file: Path, name: int):
return zstd.loads(zipfile.ZipFile(file).open(f'{name}.zst').read()) return zstd.loads(zipfile.ZipFile(file).open(f'{name}.zst').read())
def zip_to_bytes2(file: Path, name: str):
return zstd.loads(zipfile.ZipFile(file).open(name).read())
if __name__ == '__main__': if __name__ == '__main__':
pass to_zip()
# pass
# cli() # cli()
# folder_to_zip(Path('/mnt/16T/private_dataset/us/R1,U=30,M=PWI,S=(256 1502)')) # folder_to_zip(Path('/mnt/16T/private_dataset/us/R1,U=30,M=PWI,S=(256 1502)'))
# r = get(Path('/mnt/16T/private_dataset/New Folder/ST,U=60,M=FMC,S=(256 256 3002).zip'), 3) # r = get(Path('/mnt/16T/private_dataset/New Folder/ST,U=60,M=FMC,S=(256 256 3002).zip'), 3)

58
us.ipynb
View File

@ -3,8 +3,8 @@
{ {
"metadata": { "metadata": {
"ExecuteTime": { "ExecuteTime": {
"end_time": "2025-04-13T05:31:00.380352Z", "end_time": "2025-04-15T16:57:56.708615Z",
"start_time": "2025-04-13T05:31:00.366988Z" "start_time": "2025-04-15T16:57:56.374273Z"
} }
}, },
"cell_type": "code", "cell_type": "code",
@ -25,8 +25,17 @@
"from matplotlib import pyplot as plt" "from matplotlib import pyplot as plt"
], ],
"id": "5518f8a97320440e", "id": "5518f8a97320440e",
"outputs": [], "outputs": [
"execution_count": 6 {
"name": "stderr",
"output_type": "stream",
"text": [
"/home/lambda/source/scarlet/flandre/.venv/lib/python3.12/site-packages/cupyx/jit/_interface.py:173: FutureWarning: cupyx.jit.rawkernel is experimental. The interface can change in the future.\n",
" cupy._util.experimental('cupyx.jit.rawkernel')\n"
]
}
],
"execution_count": 2
}, },
{ {
"metadata": { "metadata": {
@ -137,6 +146,47 @@
} }
], ],
"execution_count": 19 "execution_count": 19
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-04-15T16:58:11.825438Z",
"start_time": "2025-04-15T16:58:00.002256Z"
}
},
"cell_type": "code",
"source": [
"from flandre.utils.RfFile import RfSequence\n",
"\n",
"rff = RfSequence.from_zip(Path('/mnt/16T/private_dataset/us/R1,U=90,M=PWI,S=(256 4502).zip'))"
],
"id": "4acfc7303fa58cf",
"outputs": [],
"execution_count": 3
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-04-15T16:59:00.287886Z",
"start_time": "2025-04-15T16:59:00.285593Z"
}
},
"cell_type": "code",
"source": "rff.frames[0].data",
"id": "75e9c5f82736241",
"outputs": [
{
"data": {
"text/plain": [
"PosixPath('/mnt/16T/private_dataset/us/R1,U=90,M=PWI,S=(256 4502).zip/0.zst')"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"execution_count": 10
} }
], ],
"metadata": { "metadata": {