flandre/src/nodes/Loader.py

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
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]
                buffer = io.BytesIO()
                buffer.write(struct.pack('=iqi', 114514, frame.meta.sequence_id, frame.meta.encoder))
                buffer.write(frame.bytes)
                playback_socket.send(buffer.getvalue())
            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