57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import io
|
|
import logging
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
import zmq
|
|
|
|
from config import PLAYBACK_SOCKET_PORT, SOFTWARE_CONFIG
|
|
from nodes.Node import Node
|
|
from utils.Msg import MoveAxisMsg, KillMsg, SetSeqMetaMsg, SeqIdMinMax, SetBaseMsg, SeqListMsg, SeqIdList, \
|
|
SetSidMsg, RfFrameWithMetaMsg
|
|
from utils.RfFile import RfSequence
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Loader(Node):
|
|
topics = [MoveAxisMsg, SetSeqMetaMsg, SetBaseMsg, SetSidMsg]
|
|
|
|
def loop(self):
|
|
playback_socket = self.context.socket(zmq.PUSH)
|
|
playback_socket.bind(f"tcp://*:{PLAYBACK_SOCKET_PORT}")
|
|
|
|
# base = Path('/mnt/16T/private_dataset/us/')
|
|
base: Path | None = None
|
|
|
|
rff = None
|
|
while True:
|
|
msg = self.recv()
|
|
if isinstance(msg, MoveAxisMsg) and rff is not None:
|
|
logger.debug(f'Move axis: {msg}')
|
|
if msg.axis == 'S':
|
|
pass
|
|
elif isinstance(msg, SetSidMsg):
|
|
frame = rff.frames[msg.value]
|
|
self.send(RfFrameWithMetaMsg(1, frame.meta, frame.bytes))
|
|
elif isinstance(msg, SetSeqMetaMsg):
|
|
if base is None:
|
|
continue
|
|
if msg.target == 'playback':
|
|
rff = RfSequence.from_folder(base / msg.name)
|
|
self.send(SeqIdMinMax(*rff.seq_id_minmax))
|
|
self.send(SeqIdList([f.meta.sequence_id for f in rff.frames]))
|
|
self.send(SetSidMsg(0))
|
|
elif isinstance(msg, SetBaseMsg):
|
|
base = Path(msg.value)
|
|
seq_list = [f.name for f in base.glob('*')]
|
|
if not seq_list:
|
|
logger.warning(f'No sequences found in {base}')
|
|
else:
|
|
SOFTWARE_CONFIG.base_dir = base
|
|
SOFTWARE_CONFIG.write_config()
|
|
self.send(SeqListMsg(seq_list))
|
|
elif isinstance(msg, KillMsg):
|
|
if msg.name == '':
|
|
break
|