From 081c1bc2f339f86dee19a7afd02e2a0747ee2840 Mon Sep 17 00:00:00 2001
From: flandre <remilia@flandre.net>
Date: Wed, 14 May 2025 00:55:31 +0800
Subject: [PATCH] add record

---
 flandre/launcher.py     | 86 ++++++++++++++++++++++++++++++++++++++++-
 flandre/nodes/Device.py | 10 ++++-
 2 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/flandre/launcher.py b/flandre/launcher.py
index 1cb9c63..827adcb 100644
--- a/flandre/launcher.py
+++ b/flandre/launcher.py
@@ -1,10 +1,12 @@
 import hashlib
 import importlib
+import json
 import logging
 import multiprocessing
 import os
 import shutil
 import subprocess
+import time
 import tomllib
 from enum import Enum, auto
 from io import TextIOWrapper
@@ -20,10 +22,12 @@ from flandre import P
 from flandre.BusClient import BusClient
 from flandre.kde_pyqt6_mainui import kde_pyqt6_mainui
 from flandre.nodes.Broker import Broker
-from flandre.nodes.Device import Device, DeviceCmd
+from flandre.nodes.Device import Device
+from flandre.nodes.Robot import Robot
 from flandre.utils.Msg import KillMsg, NodeOnlineMsg, Msg1, Msg2
 from flandre.utils.RfFrame import b2t
 from flandre.utils.mi import MiSwitch
+from flandre.utils.rtsi.serialize import DataObject
 
 
 class LaunchComponent(Enum):
@@ -211,6 +215,16 @@ def device_disable():
     dd.disable()
 
 
+@device.command('start_capture')
+def device_start_capture():
+    dd.start_capture()
+
+
+@device.command('end_capture')
+def device_end_capture():
+    dd.end_capture()
+
+
 @device.command('upload')
 @click.argument('name')
 @click.argument('file', type=click.File('r'))
@@ -285,5 +299,75 @@ def device_recvzero():
         last_b = b
 
 
+@device.command('record')
+@click.argument('folder', type=click.Path())
+def device_record(folder):
+    p = Path(folder)
+    p.mkdir(parents=True, exist_ok=True)
+    ctx = zmq.Context()
+    pull = ctx.socket(zmq.PULL)
+    pull.connect(C.live_push_socket)
+    while True:
+        b = pull.recv()
+        ts, sequence_id, encoder, buffer = b2t(b)
+        (p / f'{ts}.bin').write_bytes(b)
+
+
+robot: Robot = None
+
+
+@entrypoint.group()
+def robot():
+    global robot
+    robot = Robot()
+
+
+@robot.command('monitor')
+def robot_monitor():
+    robot.setup()
+    output1 = robot.rt.output_subscribe('actual_TCP_pose,actual_TCP_force', 250)  # 输出订阅,配方1
+    robot.rt.start()  # rtsi 开始
+    while True:
+        recv_out: DataObject = robot.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
+            print('fx', fx)
+
+
+@robot.command('record')
+@click.argument('folder', type=click.Path())
+def robot_record(folder):
+    p = Path(folder)
+    p.mkdir(parents=True, exist_ok=True)
+    robot.setup()
+    output1 = robot.rt.output_subscribe('actual_TCP_pose,actual_TCP_force', 250)  # 输出订阅,配方1
+    robot.rt.start()  # rtsi 开始
+    while True:
+        recv_out: DataObject = robot.rt.get_output_data()
+        if recv_out is None:
+            continue
+        if recv_out.recipe_id == output1.id:
+            x, y, z, rx, ry, rz = recv_out.actual_TCP_pose
+            fx, fy, fz, frx, fry, frz = recv_out.actual_TCP_force
+            d = dict(
+                x=x,
+                y=y,
+                z=z,
+                fx=fx,
+                fy=fy,
+                fz=fz,
+                rx=rx,
+                ry=ry,
+                rz=rz,
+                frx=frx,
+                fry=fry,
+                frz=frz,
+            )
+            (p / f'{time.time_ns()}.rtsi.json').write_text(json.dumps(d))
+
+
 if __name__ == '__main__':
     entrypoint()
diff --git a/flandre/nodes/Device.py b/flandre/nodes/Device.py
index 5aad4b4..8b4de37 100644
--- a/flandre/nodes/Device.py
+++ b/flandre/nodes/Device.py
@@ -34,6 +34,9 @@ class DeviceCmd(Enum):
     SetEnableOff = auto()
     SetZero = auto()
 
+    SetStartCapture = auto()
+    SetEndCapture = auto()
+
 
 class Device(Node):
     magic = 7355608
@@ -172,12 +175,17 @@ class Device(Node):
             return False
 
     def get_data(self):
-
         return self.device_cmd(DeviceCmd.GetData)
 
     def set_zero(self):
         return self.device_cmd(DeviceCmd.SetZero)
 
+    def start_capture(self):
+        return self.device_cmd(DeviceCmd.SetStartCapture)
+
+    def end_capture(self):
+        return self.device_cmd(DeviceCmd.SetEndCapture)
+
     def custom_setup(self):
         self.rep_socket = self.context.socket(zmq.REP)
         self.rep_socket.bind(f"tcp://localhost:{C.driver_rep_port}")