From ca7c8f71dfb9b223039c2c12f32ea020dad45adb Mon Sep 17 00:00:00 2001
From: flandre <remilia@flandre.net>
Date: Sun, 19 Jan 2025 18:49:41 +0800
Subject: [PATCH] add record

---
 src/nodes/Beamformer.py | 18 ++++++++++++++++--
 src/nodes/MainUI.py     |  9 ++++++---
 src/utils/Msg.py        |  7 +++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/nodes/Beamformer.py b/src/nodes/Beamformer.py
index 23fb64a..7e95fc0 100644
--- a/src/nodes/Beamformer.py
+++ b/src/nodes/Beamformer.py
@@ -1,5 +1,6 @@
 import logging
 import struct
+from pathlib import Path
 
 import cupy as cp
 import cv2
@@ -8,7 +9,7 @@ import zmq
 
 from config import PLAYBACK_SOCKET, VIDEO_WIDTH, VIDEO_HEIGHT, LIVE_SOCKET
 from nodes.Node import Node
-from utils.Msg import BMMsg, ImageArgMsg, KillMsg, SeqMetaMsg, SetPlayMode, SetDeviceConfigMsg
+from utils.Msg import BMMsg, ImageArgMsg, KillMsg, SeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg
 from utils.RfFile import RfFrame, RfSequenceMeta
 from utils.RfMat import RfMat
 
@@ -16,12 +17,15 @@ logger = logging.getLogger(__name__)
 
 
 class Beamformer(Node):
-    topics = [ImageArgMsg, SeqMetaMsg, SetPlayMode, SetDeviceConfigMsg]
+    topics = [ImageArgMsg, SeqMetaMsg, SetPlayMode, SetDeviceConfigMsg, SetRecordMsg]
 
     def __init__(self):
         super(Beamformer, self).__init__()
         self.arg = ImageArgMsg('', v1=1499)
         self.seq_meta = None
+        self.record_enable = False
+        self.record_commit = None
+        self.record_path: Path | None = None
 
     def process(self, data: RfMat):
         if data is None:
@@ -62,6 +66,8 @@ class Beamformer(Node):
                             encoder=encoder,
                             sequence_id=sequence_id,
                         ), self.seq_meta)
+                        if self.record_enable:
+                            (self.record_path / mat.frame_meta.filename).write_bytes(mat.__bytes__())
 
             if self.c.sub in socks and socks[self.c.sub] == zmq.POLLIN:
                 msg = self.recv()
@@ -79,5 +85,13 @@ class Beamformer(Node):
                     elif msg.value == 'playback':
                         device_socket.disconnect(f"tcp://{LIVE_SOCKET}")
                         device_socket.connect(f"tcp://{PLAYBACK_SOCKET}")
+                if isinstance(msg, SetRecordMsg):
+                    self.record_enable = msg.enable
+                    if msg.enable:
+                        self.record_commit = msg.commit
+                        self.seq_meta.commit = msg.commit
+                        p = Path('/tmp/flandre_record') / self.seq_meta.name
+                        p.mkdir(parents=True, exist_ok=True)
+                        self.record_path = p
 
             self.process(mat)
diff --git a/src/nodes/MainUI.py b/src/nodes/MainUI.py
index 2123e90..d405ede 100644
--- a/src/nodes/MainUI.py
+++ b/src/nodes/MainUI.py
@@ -8,7 +8,7 @@ from PyQt6.QtWidgets import QMainWindow, QApplication, QFrame
 from Main import Ui_MainWindow
 from utils.Msg import KillMsg, Msg, ImageArgMsg, SelectSeqMsg, SeqIdMinMax, MoveAxisMsg, SeqListMsg, SetBaseMsg, \
     SeqMetaMsg, SetPlayMode, DeviceConnectedMsg, DeviceEnabledMsg, DeviceOnlineMsg, SetDeviceEnabledMsg, \
-    SetDeviceConnectedMsg, DeviceConfigListMsg, SetDeviceConfigMsg
+    SetDeviceConnectedMsg, DeviceConfigListMsg, SetDeviceConfigMsg, SetRecordMsg
 from ZMQReceiver import ZMQReceiver
 from nodes.Node import Node
 from utils.RfFile import RfSequenceMeta
@@ -59,12 +59,15 @@ class Adv(QMainWindow, Ui_MainWindow):
 
     def on_record(self):
         if self.record:
+            self.p.send(SetRecordMsg(False))
             self.record = False
             self.b_record.setStyleSheet('')
 
         else:
-            self.record = True
-            self.b_record.setStyleSheet('background-color: red;')
+            if self.l_record_commit.text() != '':
+                self.p.send(SetRecordMsg(True, self.l_record_commit.text()))
+                self.record = True
+                self.b_record.setStyleSheet('background-color: red;')
 
     def on_zmq_event(self, msg: QByteArray):
         msg = Msg.decode_msg(msg.data())
diff --git a/src/utils/Msg.py b/src/utils/Msg.py
index 0acd471..304c697 100644
--- a/src/utils/Msg.py
+++ b/src/utils/Msg.py
@@ -26,6 +26,7 @@ class BG(Enum):
     SetDeviceConfigMsg = auto()
     DeviceOnlineMsg = auto()
     DeviceConfigListMsg = auto()
+    SetRecordMsg = auto()
 
 
 class Msg:
@@ -110,6 +111,12 @@ class DeviceOnlineMsg(BoolMsg):
     pass
 
 
+@dataclasses.dataclass
+class SetRecordMsg(Msg):
+    enable: bool
+    commit: str = ''
+
+
 class SetDeviceConfigMsg(StrMsg):
     pass