add recorder, fix robot arg, record robot arg

src/nodes/Loader.py

@@ -33,10 +33,6 @@ class Loader(Node):
                     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())
                 self.send(RfFrameWithMetaMsg(1, frame.meta, frame.bytes))
             elif isinstance(msg, SetSeqMetaMsg):
                 if base is None:

src/nodes/MainUI.py

@@ -329,8 +329,8 @@ class Adv(QMainWindow, Ui_MainWindow):
                 for name, txt in msg.arr:
                     self.c_live_seq_name.addItem(name, txt)
             elif isinstance(msg, RobotRtsiMsg):
-                self.l_robot_pos.setText(','.join([f'{v:.3f}' for v in msg.pos]))
+                self.l_robot_pos.setText(','.join([f'{v}' for v in msg.pos]))
-                self.l_robot_force.setText(','.join([f'{v:.3f}' for v in msg.force]))
+                self.l_robot_force.setText(','.join([f'{v}' for v in msg.force]))
             elif isinstance(msg, RecordFrameMsg):
 
                 self.record_frame_cnt += 1

src/nodes/Muxer.py

@@ -7,25 +7,22 @@ import zmq
 
 from config import IMAGING_CONFIG
 from nodes.Node import Node
-from utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, \
+from utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, \
-    ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RequestRfFrameMsg
+    ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RequestRfFrameMsg, RobotRtsiMsg
 from utils.RfFile import RfSequenceMeta
 
 logger = logging.getLogger(__name__)
 
 
 class Muxer(Node):
-    topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, RfFrameWithMetaMsg, RequestRfFrameMsg,
+    topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, RfFrameWithMetaMsg, RequestRfFrameMsg,
-              ImageArgMsg]
+              ImageArgMsg, RobotRtsiMsg]
 
     def __init__(self, level=logging.INFO):
         super(Muxer, self).__init__(level=level)
         self.seq_meta = None
         self.seq_meta_live = None
         self.seq_meta_playback = None
-        self.record_enable = False
-        self.record_commit = None
-        self.record_path: Path | None = None
         self.play_mode = None
 
     @property
@@ -37,6 +34,10 @@ class Muxer(Node):
         return seq_meta
 
     def loop(self):
+        self.rtsi = RobotRtsiMsg(
+            pos=(0, 0, 0, 0, 0, 0),
+            force=(0, 0, 0, 0, 0, 0),
+        )
         device_socket = self.context.socket(zmq.PULL)
         self.sf = None
         self.arg = ImageArgMsg('', t_start=0, t_end=1499)
@@ -52,6 +53,18 @@ class Muxer(Node):
                     msg = self.recv()
                     if isinstance(msg, RfFrameWithMetaMsg):
                         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() +
@@ -87,12 +100,5 @@ class Muxer(Node):
                     elif isinstance(msg, SetPlayMode):
                         logger.debug(f'set playmode {msg}')
                         self.play_mode = msg.value
-                    elif isinstance(msg, SetRecordMsg):
+                    elif isinstance(msg, RobotRtsiMsg):
-                        self.record_enable = msg.enable
+                        self.rtsi = msg
-                        if msg.enable:
-                            seq_meta = self.current_seq_meta
-                            self.record_commit = msg.commit
-                            seq_meta.commit = msg.commit
-                            p = Path(msg.base) / seq_meta.name
-                            p.mkdir(parents=True, exist_ok=True)
-                            self.record_path = p

src/nodes/Recorder.py

@@ -0,0 +1,87 @@
+import logging
+import struct
+import time
+from pathlib import Path
+
+import numpy as np
+import zmq
+
+from config import IMAGING_CONFIG, LIVE_SOCKET
+from nodes.Node import Node
+from utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, \
+    ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RequestRfFrameMsg, RecordFrameMsg, RobotRtsiMsg
+from utils.RfFile import RfSequenceMeta
+from utils.RfMat import RfMat
+from utils.RfMeta import RfFrameMeta
+
+logger = logging.getLogger(__name__)
+
+
+class Recorder(Node):
+    topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, RfFrameWithMetaMsg, RequestRfFrameMsg,
+              ImageArgMsg, RobotRtsiMsg]
+
+    def __init__(self, level=logging.INFO):
+        super(Recorder, self).__init__(level=level)
+        self.seq_meta_live = None
+        self.record_enable = False
+        self.record_commit = None
+        self.record_path: Path | None = None
+        self.play_mode = None
+
+    def loop(self):
+        self.rtsi = RobotRtsiMsg(
+            pos=(0, 0, 0, 0, 0, 0),
+            force=(0, 0, 0, 0, 0, 0),
+        )
+        device_socket = self.context.socket(zmq.PULL)
+        self.c.poller.register(device_socket, zmq.POLLIN)
+        time.sleep(1)
+        while True:
+            socks = dict(self.c.poller.poll())
+            for k in socks:
+                if k == device_socket:
+                    buffer = device_socket.recv()
+                    _, sequence_id, encoder = struct.unpack_from('=iqi', buffer)
+                    s = buffer[4 + 8 + 4:]
+                    if self.seq_meta_live is not None and s.__len__() // 2 == np.prod(self.seq_meta_live.shape):
+                        (self.record_path / RfFrameMeta(
+                            encoder=encoder,
+                            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],
+                        ).filename).write_bytes(s)
+                        self.send(RecordFrameMsg(s.__len__(), sequence_id))
+                if k == self.c.sub:
+                    msg = self.recv()
+                    if isinstance(msg, KillMsg):
+                        if msg.name == '':
+                            return
+                    elif isinstance(msg, SetSeqMetaMsg):
+                        match msg.target:
+                            case 'live':
+                                self.seq_meta_live = RfSequenceMeta.from_name(msg.name)
+                    elif isinstance(msg, SetRecordMsg):
+                        self.record_enable = msg.enable
+                        if msg.enable:
+                            seq_meta = self.seq_meta_live
+                            self.record_commit = msg.commit
+                            seq_meta.commit = msg.commit
+                            p = Path(msg.base) / seq_meta.name
+                            p.mkdir(parents=True, exist_ok=True)
+                            self.record_path = p
+                            device_socket.connect(f"tcp://{LIVE_SOCKET}")
+                        else:
+                            device_socket.disconnect(f"tcp://{LIVE_SOCKET}")
+                    elif isinstance(msg, RobotRtsiMsg):
+                        self.rtsi = msg

src/nodes/Robot.py

@@ -31,16 +31,32 @@ class Robot(Node):
     def rtsi_thread(self):
         rtsi_push_socket = self.context.socket(zmq.PUSH)
         rtsi_push_socket.bind('inproc://rtsi')
-        output1 = self.rt.output_subscribe('actual_joint_positions,actual_TCP_force', 125)  # 输出订阅,配方1
+        output1 = self.rt.output_subscribe('actual_TCP_pose,actual_TCP_force', 125)  # 输出订阅,配方1
         self.rt.start()  # rtsi 开始
         while not self.rtsi_thread_stop:
             recv_out: DataObject = self.rt.get_output_data()
             if recv_out is None:
                 continue
             if recv_out.recipe_id == output1.id:
+                x, y, z, r, p, yy = recv_out.actual_TCP_pose
+                fx, fy, fz, fr, fp, fyy = recv_out.actual_TCP_force
                 rtsi_push_socket.send(RobotRtsiMsg(
-                    pos=recv_out.actual_joint_positions,
+                    pos=(
-                    force=recv_out.actual_TCP_force,
+                        int(x * 100000),
+                        int(y * 100000),
+                        int(z * 100000),
+                        int(r * 1000),
+                        int(p * 1000),
+                        int(yy * 1000),
+                    ),
+                    force=(
+                        int(fx * 10),
+                        int(fy * 10),
+                        int(fz * 10),
+                        int(fr * 100),
+                        int(fp * 100),
+                        int(fyy * 100),
+                    ),
                 ).encode_msg())
         self.rt.disconnect()
 

src/utils/Msg.py

@@ -310,8 +310,8 @@ class BeamformerMsg(BytesMsg):
 
 @dataclasses.dataclass
 class RobotRtsiMsg(Msg):
-    pos: tuple[float, float, float, float, float, float]
+    pos: tuple[int, int, int, int, int, int]
-    force: tuple[float, float, float, float, float, float]
+    force: tuple[int, int, int, int, int, int]
 
 
 class RequestRfFrameMsg(Msg):

src/utils/RfMeta.py

@@ -90,6 +90,18 @@ class RfFrameMeta(RfMeta):
     robot_y: Annotated[int, 'Y'] = None
     robot_z: Annotated[int, 'Z'] = None
 
+    robot_roll: Annotated[int, 'RX'] = None
+    robot_pitch: Annotated[int, 'RY'] = None
+    robot_yal: Annotated[int, 'RZ'] = None
+
+    robot_force_x: Annotated[int, 'FX'] = None
+    robot_force_y: Annotated[int, 'FY'] = None
+    robot_force_z: Annotated[int, 'FZ'] = None
+
+    robot_force_roll: Annotated[int, 'FRX'] = None
+    robot_force_pitch: Annotated[int, 'FRY'] = None
+    robot_force_yal: Annotated[int, 'FRZ'] = None
+
 
 @dataclass
 class RfSequenceMeta(RfMeta):

test/kdemain.py

@@ -9,6 +9,7 @@ from nodes.ImageCV import ImageCV
 from nodes.ImageFFMPEG import ImageFFMPEG
 from nodes.Loader import Loader
 from nodes.Muxer import Muxer
+from nodes.Recorder import Recorder
 from nodes.Robot import Robot
 from qtonly import kde_pyqt6_mainui
 from utils.Msg import KillMsg
@@ -28,6 +29,7 @@ if __name__ == '__main__':
         Loader(),
         Muxer(level=logging.DEBUG),
         Robot(),
+        Recorder(),
     ]
     for p in ps:
         pps.append(multiprocessing.Process(target=p))