Merge 61202761f5 into 768e36660d

lerobot/common/datasets/lerobot_dataset.py

@@ -967,17 +967,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
         Starts recording audio data provided by the microphone and directly writes it in a .wav file.
         """
 
-        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
-            )
-        )
+        audio_file = self._get_raw_audio_file_path(self.num_episodes, "observation.audio." + microphone_key)
+        microphone.start_recording(output_file=audio_file)
 
     def save_episode(self, episode_data: dict | None = None) -> None:
         """
@@ -1075,11 +1066,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             shutil.rmtree(self.root / "images")
 
         # delete raw audio files
-        raw_audio_files = list(self.root.rglob("*.wav"))
-        for raw_audio_file in raw_audio_files:
-            raw_audio_file.unlink()
-            if len(list(raw_audio_file.parent.iterdir())) == 0:
-                raw_audio_file.parent.rmdir()
+        img_dir = self.root / "raw_audio"
+        if img_dir.is_dir():
+            shutil.rmtree(self.root / "raw_audio")
 
         if not episode_data:  # Reset the buffer
             self.episode_buffer = self.create_episode_buffer()

lerobot/common/datasets/utils.py

@@ -56,7 +56,7 @@ TASKS_PATH = "meta/tasks.jsonl"
 DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
 DEFAULT_PARQUET_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
 DEFAULT_IMAGE_PATH = "images/{image_key}/episode_{episode_index:06d}/frame_{frame_index:06d}.png"
-DEFAULT_RAW_AUDIO_PATH = "audio/{audio_key}/episode_{episode_index:06d}.wav"
+DEFAULT_RAW_AUDIO_PATH = "raw_audio/{audio_key}/episode_{episode_index:06d}.wav"
 DEFAULT_COMPRESSED_AUDIO_PATH = "audio/chunk-{episode_chunk:03d}/{audio_key}/episode_{episode_index:06d}.m4a"
 
 DEFAULT_AUDIO_CHUNK_DURATION = 0.5  # seconds

lerobot/common/robot_devices/microphones/microphone.py

@@ -302,7 +302,12 @@ class Microphone:
 
         return audio_readings
 
-    def start_recording(self, output_file: str | None = None, multiprocessing: bool | None = False) -> None:
+    def start_recording(
+        self,
+        output_file: str | None = None,
+        multiprocessing: bool | None = False,
+        overwrite: bool | None = True,
+    ) -> None:
         """
         Starts the recording of the microphone. If output_file is provided, the audio will be written to this file.
         """
@@ -323,8 +328,15 @@ class Microphone:
         # Write recordings into a file if output_file is provided
         if output_file is not None:
             output_file = Path(output_file)
+            output_file.parent.mkdir(parents=True, exist_ok=True)
+
             if output_file.exists():
-                output_file.unlink()
+                if overwrite:
+                    output_file.unlink()
+                else:
+                    raise FileExistsError(
+                        f"Output file {output_file} already exists. Set overwrite to True to overwrite it."
+                    )
 
             if multiprocessing:
                 self.record_stop_event = process_Event()

lerobot/common/robot_devices/microphones/utils.py

@@ -21,7 +21,12 @@ from lerobot.common.robot_devices.microphones.configs import MicrophoneConfig, M
 class Microphone(Protocol):
     def connect(self): ...
     def disconnect(self): ...
-    def start_recording(self, output_file: str | None = None, multiprocessing: bool | None = False): ...
+    def start_recording(
+        self,
+        output_file: str | None = None,
+        multiprocessing: bool | None = False,
+        overwrite: bool | None = True,
+    ): ...
     def stop_recording(self): ...