Adding microphone recording in control loop
This commit is contained in:
parent
8ee61bb81f
commit
e4eebd0680
|
@ -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.
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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:
|
||||
|
||||
|
|
Loading…
Reference in New Issue