mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
Adding microphone recording in control loop
This commit is contained in:
@@ -79,6 +79,7 @@ from lerobot.datasets.video_utils import (
|
|||||||
get_safe_default_codec,
|
get_safe_default_codec,
|
||||||
get_video_info,
|
get_video_info,
|
||||||
)
|
)
|
||||||
|
from lerobot.microphones import Microphone
|
||||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||||
|
|
||||||
CODEBASE_VERSION = "v3.0"
|
CODEBASE_VERSION = "v3.0"
|
||||||
@@ -1282,6 +1283,23 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
|
|
||||||
self.episode_buffer["size"] += 1
|
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(
|
def save_episode(
|
||||||
self,
|
self,
|
||||||
episode_data: dict | None = None,
|
episode_data: dict | None = None,
|
||||||
@@ -1780,8 +1798,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
episode_index = self.episode_buffer["episode_index"]
|
episode_index = self.episode_buffer["episode_index"]
|
||||||
if isinstance(episode_index, np.ndarray):
|
if isinstance(episode_index, np.ndarray):
|
||||||
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
||||||
for microphone_key in self.meta.microphone_keys:
|
for audio_key in self.meta.audio_keys:
|
||||||
audio_file = self.root / self.meta.get_audio_file_path(episode_index, microphone_key)
|
audio_file = self.root / self.meta.get_audio_file_path(episode_index, audio_key)
|
||||||
if audio_file.is_file():
|
if audio_file.is_file():
|
||||||
audio_file.unlink()
|
audio_file.unlink()
|
||||||
|
|
||||||
|
|||||||
@@ -229,9 +229,8 @@ class Microphone:
|
|||||||
self.thread.daemon = True
|
self.thread.daemon = True
|
||||||
self.thread.start()
|
self.thread.start()
|
||||||
|
|
||||||
self.stream.start()
|
|
||||||
|
|
||||||
self.logs["start_timestamp"] = capture_timestamp_utc()
|
self.logs["start_timestamp"] = capture_timestamp_utc()
|
||||||
|
self.stream.start()
|
||||||
|
|
||||||
def stop_recording(self) -> None:
|
def stop_recording(self) -> None:
|
||||||
if not self.is_connected:
|
if not self.is_connected:
|
||||||
|
|||||||
@@ -313,6 +313,10 @@ def record_loop(
|
|||||||
preprocessor.reset()
|
preprocessor.reset()
|
||||||
postprocessor.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
|
timestamp = 0
|
||||||
start_episode_t = time.perf_counter()
|
start_episode_t = time.perf_counter()
|
||||||
while timestamp < control_time_s:
|
while timestamp < control_time_s:
|
||||||
@@ -397,6 +401,9 @@ def record_loop(
|
|||||||
|
|
||||||
timestamp = time.perf_counter() - start_episode_t
|
timestamp = time.perf_counter() - start_episode_t
|
||||||
|
|
||||||
|
for _, microphone in robot.microphones.items():
|
||||||
|
microphone.stop_recording()
|
||||||
|
|
||||||
|
|
||||||
@parser.wrap()
|
@parser.wrap()
|
||||||
def record(cfg: RecordConfig) -> LeRobotDataset:
|
def record(cfg: RecordConfig) -> LeRobotDataset:
|
||||||
|
|||||||
Reference in New Issue
Block a user