Adding audio tests

tests/datasets/test_compute_stats.py

@@ -26,16 +26,21 @@ from lerobot.datasets.compute_stats import (
     compute_episode_stats,
     estimate_num_samples,
     get_feature_stats,
+    sample_audio,
     sample_images,
     sample_indices,
 )
-from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
+from lerobot.utils.constants import OBS_AUDIO, OBS_IMAGE, OBS_STATE
 
 
 def mock_load_image_as_numpy(path, dtype, channel_first):
     return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype)
 
 
+def mock_load_audio(path):
+    return np.ones((16000, 2), dtype=np.float32)
+
+
 @pytest.fixture
 def sample_array():
     return np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
@@ -73,6 +78,16 @@ def test_sample_images(mock_load):
     assert len(images) == estimate_num_samples(100)
 
 
+@patch("lerobot.datasets.compute_stats.load_audio", side_effect=mock_load_audio)
+def test_sample_audio(mock_load):
+    audio_path = "audio.wav"
+    audio_samples = sample_audio(audio_path)
+    assert isinstance(audio_samples, np.ndarray)
+    assert audio_samples.shape[1] == 2
+    assert audio_samples.dtype == np.float32
+    assert len(audio_samples) == estimate_num_samples(16000)
+
+
 def test_get_feature_stats_images():
     data = np.random.rand(100, 3, 32, 32)
     stats = get_feature_stats(data, axis=(0, 2, 3), keepdims=True)
@@ -81,6 +96,14 @@ def test_get_feature_stats_images():
     assert stats["min"].shape == stats["max"].shape == stats["mean"].shape == stats["std"].shape
 
 
+def test_get_feature_stats_audio():
+    data = np.random.uniform(-1, 1, (16000, 2))
+    stats = get_feature_stats(data, axis=0, keepdims=True)
+    assert "min" in stats and "max" in stats and "mean" in stats and "std" in stats and "count" in stats
+    np.testing.assert_equal(stats["count"], np.array([16000]))
+    assert stats["min"].shape == stats["max"].shape == stats["mean"].shape == stats["std"].shape
+
+
 def test_get_feature_stats_axis_0_keepdims(sample_array):
     expected = {
         "min": np.array([[1, 2, 3]]),
@@ -145,20 +168,27 @@ def test_get_feature_stats_single_value():
 def test_compute_episode_stats():
     episode_data = {
         OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)],
+        OBS_AUDIO: "audio.wav",
         OBS_STATE: np.random.rand(100, 10),
     }
     features = {
         OBS_IMAGE: {"dtype": "image"},
+        OBS_AUDIO: {"dtype": "audio"},
         OBS_STATE: {"dtype": "numeric"},
     }
 
-    with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
+    with (
+        patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy),
+        patch("lerobot.datasets.compute_stats.load_audio", side_effect=mock_load_audio),
+    ):
         stats = compute_episode_stats(episode_data, features)
 
-    assert OBS_IMAGE in stats and OBS_STATE in stats
-    assert stats[OBS_IMAGE]["count"].item() == 100
-    assert stats[OBS_STATE]["count"].item() == 100
+    assert OBS_IMAGE in stats and OBS_AUDIO in stats and OBS_STATE in stats
+    assert stats[OBS_IMAGE]["count"].item() == estimate_num_samples(100)
+    assert stats[OBS_AUDIO]["count"].item() == estimate_num_samples(16000)
+    assert stats[OBS_STATE]["count"].item() == estimate_num_samples(100)
     assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1)
+    assert stats[OBS_AUDIO]["mean"].shape == (1, 2)
 
 
 def test_assert_type_and_shape_valid():

tests/datasets/test_datasets.py

@@ -37,6 +37,7 @@ from lerobot.datasets.lerobot_dataset import (
     _encode_video_worker,
 )
 from lerobot.datasets.utils import (
+    DEFAULT_AUDIO_CHUNK_DURATION,
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_FILE_SIZE_IN_MB,
@@ -49,7 +50,13 @@ from lerobot.envs.factory import make_env_config
 from lerobot.policies.factory import make_policy_config
 from lerobot.robots import make_robot_from_config
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD
-from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
+from tests.fixtures.constants import (
+    DEFAULT_SAMPLE_RATE,
+    DUMMY_AUDIO_CHANNELS,
+    DUMMY_CHW,
+    DUMMY_HWC,
+    DUMMY_REPO_ID,
+)
 from tests.mocks.mock_robot import MockRobotConfig
 from tests.utils import require_x86_64_kernel
 
@@ -70,6 +77,20 @@ def image_dataset(tmp_path, empty_lerobot_dataset_factory):
     return empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
 
 
+@pytest.fixture
+def audio_dataset(tmp_path, empty_lerobot_dataset_factory):
+    features = {
+        "audio": {
+            "dtype": "audio",
+            "shape": (DUMMY_AUDIO_CHANNELS,),
+            "names": [
+                "channels",
+            ],
+        }
+    }
+    return empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
+
+
 def test_same_attributes_defined(tmp_path, lerobot_dataset_factory):
     """
     Instantiate a LeRobotDataset both ways with '__init__()' and 'create()' and verify that instantiated
@@ -411,6 +432,23 @@ def test_tmp_mixed_deletion(tmp_path, empty_lerobot_dataset_factory):
     )
 
 
+def test_add_frame_audio(audio_dataset):
+    dataset = audio_dataset
+    dataset.add_frame(
+        {
+            "audio": np.random.rand(
+                int(DEFAULT_AUDIO_CHUNK_DURATION * DEFAULT_SAMPLE_RATE), DUMMY_AUDIO_CHANNELS
+            )
+        },
+        task="Dummy task",
+    )
+    dataset.save_episode()
+
+    assert dataset[0]["audio"].shape == torch.Size(
+        (int(DEFAULT_AUDIO_CHUNK_DURATION * DEFAULT_SAMPLE_RATE), DUMMY_AUDIO_CHANNELS)
+    )
+
+
 # TODO(aliberts):
 # - [ ] test various attributes & state from init and create
 # - [ ] test init with episodes and check num_frames
@@ -450,6 +488,7 @@ def test_factory(env_name, repo_id, policy_name):
     dataset = make_dataset(cfg)
     delta_timestamps = dataset.delta_timestamps
     camera_keys = dataset.meta.camera_keys
+    audio_keys = dataset.meta.audio_keys
 
     item = dataset[0]
 
@@ -492,6 +531,11 @@ def test_factory(env_name, repo_id, policy_name):
                 # test c,h,w
                 assert item[key].shape[0] == 3, f"{key}"
 
+        for key in audio_keys:
+            assert item[key].dtype == torch.float32, f"{key}"
+            assert item[key].max() <= 1.0, f"{key}"
+            assert item[key].min() >= -1.0, f"{key}"
+
     if delta_timestamps is not None:
         # test missing keys in delta_timestamps
         for key in delta_timestamps: