Adding microphone recording in control loop

This commit is contained in:
CarolinePascal 2025-03-28 17:16:17 +01:00
parent 8ee61bb81f
commit e4eebd0680
No known key found for this signature in database
3 changed files with 26 additions and 2 deletions

View File

@ -79,6 +79,7 @@ from lerobot.common.datasets.video_utils import (
get_audio_info,
)
from lerobot.common.robot_devices.robots.utils import Robot
from lerobot.common.robot_devices.microphones.utils import Microphone
CODEBASE_VERSION = "v2.1"
@ -910,6 +911,17 @@ 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) -> None:
"""
This will save to disk the current episode in self.episode_buffer.

View File

@ -78,6 +78,11 @@ def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, f
if key in robot.logs:
log_dt(f"dtR{name}", robot.logs[key])
for name in robot.microphones:
key = f"read_microphone_{name}_dt_s"
if key in robot.logs:
log_dt(f"dtR{name}", robot.logs[key])
info_str = " ".join(log_items)
logging.info(info_str)
@ -243,6 +248,11 @@ def control_loop(
timestamp = 0
start_episode_t = time.perf_counter()
if teleoperate and dataset is not None:
for microphone_key, microphone in robot.microphones.items():
dataset.add_microphone_recording(microphone, microphone_key)
while timestamp < control_time_s:
start_loop_t = time.perf_counter()
@ -286,6 +296,9 @@ def control_loop(
events["exit_early"] = False
break
for _, microphone in robot.microphones.items():
microphone.stop_recording()
def reset_environment(robot, events, reset_time_s, fps):
# TODO(rcadene): refactor warmup_record and reset_environment

View File

@ -238,9 +238,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: