flandre/src/nodes/Loader.py

import io
import logging
import struct
from pathlib import Path

import zmq

from config import PLAYBACK_SOCKET_PORT
from nodes.Node import Node
from utils.Msg import MoveAxisMsg, KillMsg, SelectSeqMsg, SeqMetaMsg, SeqIdMinMax, SetBaseMsg, SeqListMsg
from utils.RfFile import RfSequence

logger = logging.getLogger(__name__)


class Loader(Node):
    topics = [MoveAxisMsg, SelectSeqMsg, SetBaseMsg]

    def loop(self):
        s2 = self.context.socket(zmq.PUSH)
        s2.bind(f"tcp://*:{PLAYBACK_SOCKET_PORT}")

        # base = Path('/mnt/16T/private_dataset/us/')
        base: Path | None = None

        rff = None
        while True:
            r = self.recv()
            if isinstance(r, MoveAxisMsg) and rff is not None:
                logger.debug(f'Move axis: {rff}')
                for frame in rff.frames:
                    if frame.meta.sequence_id == r.value:
                        buffer = io.BytesIO()
                        buffer.write(struct.pack('=iqi', 114514, frame.meta.sequence_id, frame.meta.encoder))
                        buffer.write(frame.bytes)
                        s2.send(buffer.getvalue())
            elif isinstance(r, SelectSeqMsg):
                if base is None:
                    continue
                rff = RfSequence.from_folder(base / r.value)
                self.send(SeqMetaMsg(rff.meta.name))
                self.send(SeqIdMinMax(*rff.seq_id_minmax))
            elif isinstance(r, SetBaseMsg):
                base = Path(r.value)
                self.send(SeqListMsg([f.name for f in base.glob('*')]))
            elif isinstance(r, KillMsg):
                if r.name == '':
                    break