diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 751fe017e..855398689 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/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() diff --git a/src/lerobot/microphones/microphone.py b/src/lerobot/microphones/microphone.py index cb2fbf5ce..9e0942f42 100644 --- a/src/lerobot/microphones/microphone.py +++ b/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: diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index f03776989..ac6f611f5 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/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: