yukari/flandre
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add recorder, fix robot arg, record robot arg

Browse Source
This commit is contained in:
remilia 2025-02-26 22:51:22 +08:00
parent d92a9778d4
commit 1738ae5b48
8 changed files with 146 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/nodes/Loader.py
View File

@ -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:

4
src/nodes/MainUI.py
View File

@ -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

38
src/nodes/Muxer.py
View File

@ -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

87
src/nodes/Recorder.py Normal file
View File

@ -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

22
src/nodes/Robot.py
View File

@ -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()

4
src/utils/Msg.py
View File

@ -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):

12
src/utils/RfMeta.py
View File

@ -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):

2
test/kdemain.py
View File

@ -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))