diff --git a/src/nodes/Loader.py b/src/nodes/Loader.py index 62a9281..dc9aef3 100644 --- a/src/nodes/Loader.py +++ b/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: diff --git a/src/nodes/MainUI.py b/src/nodes/MainUI.py index d70ceeb..160719d 100644 --- a/src/nodes/MainUI.py +++ b/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_force.setText(','.join([f'{v:.3f}' for v in msg.force])) + self.l_robot_pos.setText(','.join([f'{v}' for v in msg.pos])) + self.l_robot_force.setText(','.join([f'{v}' for v in msg.force])) elif isinstance(msg, RecordFrameMsg): self.record_frame_cnt += 1 diff --git a/src/nodes/Muxer.py b/src/nodes/Muxer.py index 2d88bc4..cdb9876 100644 --- a/src/nodes/Muxer.py +++ b/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, \ - ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RequestRfFrameMsg +from utils.Msg import ImageArgMsg, KillMsg, SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, \ + ImagingConfigNameListMsg, RfFrameWithMetaMsg, BeamformerMsg, RequestRfFrameMsg, RobotRtsiMsg from utils.RfFile import RfSequenceMeta logger = logging.getLogger(__name__) class Muxer(Node): - topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg, RfFrameWithMetaMsg, RequestRfFrameMsg, - ImageArgMsg] + topics = [SetSeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, RfFrameWithMetaMsg, RequestRfFrameMsg, + 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): - self.record_enable = msg.enable - 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 + elif isinstance(msg, RobotRtsiMsg): + self.rtsi = msg diff --git a/src/nodes/Recorder.py b/src/nodes/Recorder.py new file mode 100644 index 0000000..8609b2d --- /dev/null +++ b/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 diff --git a/src/nodes/Robot.py b/src/nodes/Robot.py index 64e5fc9..8c5e069 100644 --- a/src/nodes/Robot.py +++ b/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, - force=recv_out.actual_TCP_force, + pos=( + 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() diff --git a/src/utils/Msg.py b/src/utils/Msg.py index ad38a73..9d8e376 100644 --- a/src/utils/Msg.py +++ b/src/utils/Msg.py @@ -310,8 +310,8 @@ class BeamformerMsg(BytesMsg): @dataclasses.dataclass class RobotRtsiMsg(Msg): - pos: tuple[float, float, float, float, float, float] - force: tuple[float, float, float, float, float, float] + pos: tuple[int, int, int, int, int, int] + force: tuple[int, int, int, int, int, int] class RequestRfFrameMsg(Msg): diff --git a/src/utils/RfMeta.py b/src/utils/RfMeta.py index 5fb48b2..be6d9dc 100644 --- a/src/utils/RfMeta.py +++ b/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): diff --git a/test/kdemain.py b/test/kdemain.py index f7bacc6..f93cf72 100644 --- a/test/kdemain.py +++ b/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))