Adding microphone recording in control loop

src/lerobot/datasets/lerobot_dataset.py

@@ -79,6 +79,7 @@ from lerobot.datasets.video_utils import (
     get_safe_default_codec,
     get_video_info,
 )
+from lerobot.microphones import Microphone
 from lerobot.utils.constants import HF_LEROBOT_HOME
 
 CODEBASE_VERSION = "v3.0"
@@ -1282,6 +1283,23 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         self.episode_buffer["size"] += 1
 
+    def add_microphone_recording(self, microphone: Microphone, microphone_key: str) -> None:
+        """
+        This function will start recording audio from the microphone and save it to disk.
+        """
+
+        audio_dir = self._get_raw_audio_file_path(
+            self.num_episodes, "observation.audio." + microphone_key
+        ).parent
+        if not audio_dir.is_dir():
+            audio_dir.mkdir(parents=True, exist_ok=True)
+
+        microphone.start_recording(
+            output_file=self._get_raw_audio_file_path(
+                self.num_episodes, "observation.audio." + microphone_key
+            )
+        )
+
     def save_episode(
         self,
         episode_data: dict | None = None,
@@ -1780,8 +1798,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
             episode_index = self.episode_buffer["episode_index"]
             if isinstance(episode_index, np.ndarray):
                 episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
-            for microphone_key in self.meta.microphone_keys:
-                audio_file = self.root / self.meta.get_audio_file_path(episode_index, microphone_key)
+            for audio_key in self.meta.audio_keys:
+                audio_file = self.root / self.meta.get_audio_file_path(episode_index, audio_key)
                 if audio_file.is_file():
                     audio_file.unlink()
 

src/lerobot/microphones/microphone.py

@@ -229,9 +229,8 @@ class Microphone:
             self.thread.daemon = True
             self.thread.start()
 
-        self.stream.start()
-
         self.logs["start_timestamp"] = capture_timestamp_utc()
+        self.stream.start()
 
     def stop_recording(self) -> None:
         if not self.is_connected:

src/lerobot/scripts/lerobot_record.py

@@ -313,6 +313,10 @@ def record_loop(
         preprocessor.reset()
         postprocessor.reset()
 
+    if dataset is not None:
+        for microphone_key, microphone in robot.microphones.items():
+            dataset.add_microphone_recording(microphone, microphone_key)
+
     timestamp = 0
     start_episode_t = time.perf_counter()
     while timestamp < control_time_s:
@@ -397,6 +401,9 @@ def record_loop(
 
         timestamp = time.perf_counter() - start_episode_t
 
+    for _, microphone in robot.microphones.items():
+        microphone.stop_recording()
+
 
 @parser.wrap()
 def record(cfg: RecordConfig) -> LeRobotDataset: