Added CSV writing to IMU Experiment

examples/imu_experiment/imu_experiment.py

@@ -82,6 +82,6 @@ def set_setup_task(
     builder = (
         ImuExperimentBuilder()
         .set_query_task(QueryTaskWidgetBuilder().set_kafka_manager(manager))
-        .set_setup_task(SetupTaskWidgetBuilder())
+        .set_setup_task(SetupTaskWidgetBuilder().set_kafka_manager(manager))
     )
     build_and_launch_window_single_process(builder)

examples/imu_experiment/src/widgets/setup_task_widget.py

@@ -7,6 +7,8 @@
 
 from natKit.client.gui.pyqt6.widget import Task, TaskBuilder, ExperimentBuilder
 from natKit.client.gui.pyqt6.window import ExperimentWindow
+from natKit.common.kafka import KafkaManager
+from natKit.api import ImuDataSchema
 
 from .imu_plot_widget import ImuPlotWidget
 from .imu_connection_widget import ImuConnectionWidget
@@ -16,15 +18,18 @@
 
 from typing import NoReturn
 
+from time import sleep
+
 
 class SetupTaskWidget(Task):
     """ """
 
     onNext = pyqtSignal()
     onPrev = pyqtSignal()
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, kafka_manager: KafkaManager, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.kafka_manager = kafka_manager
         self.imu_streams = ImuStreams()
 
     def set_imu_streams(self, imu_streams) -> NoReturn:
@@ -59,6 +64,12 @@ def setup(self):
             self._previous_task_button, alignment=Qt.AlignmentFlag.AlignLeft
         )
 
+        self.write_to_csv_button = QPushButton(text="Write CSV")
+        self.write_to_csv_button.clicked.connect(self.write_csv)
+        button_layout.addWidget(
+            self.write_to_csv_button, alignment=Qt.AlignmentFlag.AlignCenter
+        )
+
         self._next_task_button = QPushButton(text="Data Collection ->")
         self._next_task_button.clicked.connect(lambda: self.onNext.emit())
         button_layout.addWidget(
@@ -69,13 +80,33 @@ def setup(self):
 
         self.setLayout(self.layout)
 
+    def write_csv(self):
+        streams_and_names = self.imu_streams.get_streams_and_names()
+        print("Writing to CSV...")
+        for stream, name in streams_and_names:
+            new_stream = self.kafka_manager.create_new_stream_from_stream(stream)
+            while not new_stream.stream_ready:
+                sleep(0.01)
+            sleep(0.1)
+
+            with open(name + ".csv", "w") as f:
+                f.write(ImuDataSchema.csv_header() + "\n")
+                while True:
+                    data = new_stream.read_data()
+                    if data is None:
+                        break
+                    f.write(data.to_csv() + "\n")
+        print("Finished writing to CSV")
+
 
 class SetupTaskWidgetBuilder(TaskBuilder):
     def __init__(self) -> NoReturn:
         super().__init__()
+        self.kafka_manager = None
 
     def build(self) -> Task:
         return SetupTaskWidget(
+            kafka_manager=self.kafka_manager,
             name=self.name,
             stages=self.stages,
             inter_stage_interval=self.inter_stage_interval,
@@ -84,3 +115,7 @@ def build(self) -> Task:
             events=self.events,
             duration_events=self.duration_events,
         )
+
+    def set_kafka_manager(self, manager: KafkaManager):
+        self.kafka_manager = manager
+        return self

natKit/api/src/python/natKit/api/schemas/imu_data_schema.py

@@ -56,3 +56,28 @@ def deserialize(encoder: Encoder, msg: bytes) -> Schema:
             assert 0, "{} does not support {} encoding".format(
                 ImuDataSchema.get_name(), encoder.get_name()
             )
+
+    @staticmethod
+    def csv_header() -> str:
+        header = [
+            "timestamp",
+            "data0",
+            "data1",
+            "data2",
+            "data3",
+            "data4",
+            "data5",
+            "data6",
+            "data7",
+            "data8",
+            "calibration",
+        ]
+        return ",".join(header)
+
+    def to_csv(self) -> str:
+        data = [0 for i in range(11)]
+        data[0] = str(self.timestamp)
+        for i, datum in enumerate(self.data):
+            data[i + 1] = str(datum)
+        data[10] = str(self.calibration)
+        return ",".join(data)

natKit/common/src/python/natKit/common/kafka/topic_manager.py

@@ -142,6 +142,11 @@ def _create_stream(self, meta_topic_string: str, data_topic_string: str):
             data_messenger=self.create_messenger(data_topic_string),
         )
 
+    def create_new_stream_from_stream(self, stream: Stream) -> Stream:
+        return self._create_stream(
+            stream.get_meta_topic_string(), stream.get_data_topic_string()
+        )
+
     def find_streams(self):
         topic_string_pairs = StreamHelper.find_topic_stream_pairs(
             self.query_topic_names()