mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
fix(online audio chunks): querying audio chunks as a shifting widow over previous audio samples to match the default audio chunk size
This commit is contained in:
@@ -69,6 +69,8 @@ from pathlib import Path
|
|||||||
from pprint import pformat
|
from pprint import pformat
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
from lerobot.cameras import ( # noqa: F401
|
from lerobot.cameras import ( # noqa: F401
|
||||||
CameraConfig, # noqa: F401
|
CameraConfig, # noqa: F401
|
||||||
)
|
)
|
||||||
@@ -81,8 +83,12 @@ from lerobot.configs.policies import PreTrainedConfig
|
|||||||
from lerobot.datasets.image_writer import safe_stop_image_writer
|
from lerobot.datasets.image_writer import safe_stop_image_writer
|
||||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||||
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
|
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
|
||||||
from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts
|
from lerobot.datasets.utils import DEFAULT_AUDIO_CHUNK_DURATION, build_dataset_frame, combine_feature_dicts
|
||||||
from lerobot.datasets.video_utils import VideoEncodingManager
|
from lerobot.datasets.video_utils import VideoEncodingManager
|
||||||
|
from lerobot.microphones.utils import (
|
||||||
|
async_microphones_start_recording,
|
||||||
|
async_microphones_stop_recording,
|
||||||
|
)
|
||||||
from lerobot.policies.factory import make_policy, make_pre_post_processors
|
from lerobot.policies.factory import make_policy, make_pre_post_processors
|
||||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||||
from lerobot.policies.utils import make_robot_action
|
from lerobot.policies.utils import make_robot_action
|
||||||
@@ -317,8 +323,15 @@ def record_loop(
|
|||||||
for microphone_key, microphone in robot.microphones.items():
|
for microphone_key, microphone in robot.microphones.items():
|
||||||
dataset.add_microphone_recording(microphone, microphone_key)
|
dataset.add_microphone_recording(microphone, microphone_key)
|
||||||
else:
|
else:
|
||||||
for _, microphone in robot.microphones.items():
|
async_microphones_start_recording(robot.microphones)
|
||||||
microphone.start_recording()
|
|
||||||
|
# Create a buffer for audio observations (shifting window of fixed size over audio samples)
|
||||||
|
audio_buffer = {
|
||||||
|
microphone_name: np.zeros(
|
||||||
|
(int(microphone.sample_rate * DEFAULT_AUDIO_CHUNK_DURATION), len(microphone.channels))
|
||||||
|
)
|
||||||
|
for microphone_name, microphone in robot.microphones.items()
|
||||||
|
}
|
||||||
|
|
||||||
timestamp = 0
|
timestamp = 0
|
||||||
start_episode_t = time.perf_counter()
|
start_episode_t = time.perf_counter()
|
||||||
@@ -340,6 +353,15 @@ def record_loop(
|
|||||||
|
|
||||||
# Get action from either policy or teleop
|
# Get action from either policy or teleop
|
||||||
if policy is not None and preprocessor is not None and postprocessor is not None:
|
if policy is not None and preprocessor is not None and postprocessor is not None:
|
||||||
|
# Transform instantaneous audio samples into a buffer of fixed size
|
||||||
|
for name in audio_buffer:
|
||||||
|
# Remove as many old audio samples as needed
|
||||||
|
audio_buffer[name] = audio_buffer[name][len(observation_frame[name]) :]
|
||||||
|
# Add new audio samples
|
||||||
|
audio_buffer[name] = np.vstack((audio_buffer[name], observation_frame[name]))
|
||||||
|
# Add the audio buffer to the observation
|
||||||
|
observation_frame[name] = audio_buffer[name]
|
||||||
|
|
||||||
action_values = predict_action(
|
action_values = predict_action(
|
||||||
observation=observation_frame,
|
observation=observation_frame,
|
||||||
policy=policy,
|
policy=policy,
|
||||||
@@ -404,8 +426,7 @@ def record_loop(
|
|||||||
|
|
||||||
timestamp = time.perf_counter() - start_episode_t
|
timestamp = time.perf_counter() - start_episode_t
|
||||||
|
|
||||||
for _, microphone in robot.microphones.items():
|
async_microphones_stop_recording(robot.microphones)
|
||||||
microphone.stop_recording()
|
|
||||||
|
|
||||||
|
|
||||||
@parser.wrap()
|
@parser.wrap()
|
||||||
|
|||||||
Reference in New Issue
Block a user