feat(WIP): adding support for multi-process camera video capture

* fix bug in `_copy_data_with_feature_changes`

* Update src/lerobot/datasets/dataset_tools.py

Co-authored-by: Caroline Pascal <caroline8.pascal@gmail.com>
Signed-off-by: Michel Aractingi <michel.aractingi@huggingface.co>

* add missing import

---------

Signed-off-by: Michel Aractingi <michel.aractingi@huggingface.co>
Co-authored-by: Caroline Pascal <caroline8.pascal@gmail.com>

src/lerobot/cameras/opencv/camera_opencv.py

@@ -23,6 +23,8 @@ import platform
 import time
 from pathlib import Path
 from threading import Event, Lock, Thread
+from multiprocessing import Process, Event as EventProcess, JoinableQueue as Queue
+from queue import Empty
 from typing import Any
 
 from numpy.typing import NDArray  # type: ignore  # TODO: add type stubs for numpy.typing
@@ -119,11 +121,10 @@ class OpenCVCamera(Camera):
 
         self.videocapture: cv2.VideoCapture | None = None
 
-        self.thread: Thread | None = None
-        self.stop_event: Event | None = None
-        self.frame_lock: Lock = Lock()
+        self.process: Process | None = None
+        self.stop_event: EventProcess | None = None
+        self.frame_queue: Queue = Queue()
         self.latest_frame: NDArray[Any] | None = None
-        self.new_frame_event: Event = Event()
 
         self.rotation: int | None = get_cv2_rotation(config.rotation)
         self.backend: int = get_cv2_backend()
@@ -442,37 +443,36 @@ class OpenCVCamera(Camera):
         while not self.stop_event.is_set():
             try:
                 color_image = self.read()
-
-                with self.frame_lock:
-                    self.latest_frame = color_image
-                self.new_frame_event.set()
+                self.frame_queue.put_nowait(color_image)
 
             except DeviceNotConnectedError:
                 break
             except Exception as e:
                 logger.warning(f"Error reading frame in background thread for {self}: {e}")
 
-    def _start_read_thread(self) -> None:
+    def _start_read_process(self) -> None:
         """Starts or restarts the background read thread if it's not running."""
-        if self.thread is not None and self.thread.is_alive():
-            self.thread.join(timeout=0.1)
+        if self.process is not None and self.process.is_alive():
+            self.frame_queue.join()
+            self.process.join()
         if self.stop_event is not None:
             self.stop_event.set()
 
         self.stop_event = Event()
-        self.thread = Thread(target=self._read_loop, args=(), name=f"{self}_read_loop")
-        self.thread.daemon = True
-        self.thread.start()
+        self.process = Process(target=self._read_loop, args=(), name=f"{self}_read_loop")
+        self.process.daemon = True
+        self.process.start()
 
     def _stop_read_thread(self) -> None:
         """Signals the background read thread to stop and waits for it to join."""
         if self.stop_event is not None:
             self.stop_event.set()
 
-        if self.thread is not None and self.thread.is_alive():
-            self.thread.join(timeout=2.0)
+        if self.process is not None and self.process.is_alive():
+            self.frame_queue.join()
+            self.process.join()
 
-        self.thread = None
+        self.process = None
         self.stop_event = None
 
     def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
@@ -499,24 +499,32 @@ class OpenCVCamera(Camera):
         if not self.is_connected:
             raise DeviceNotConnectedError(f"{self} is not connected.")
 
-        if self.thread is None or not self.thread.is_alive():
-            self._start_read_thread()
+        if self.process is None or not self.process.is_alive():
+            self._start_read_process()
 
-        if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0):
-            thread_alive = self.thread is not None and self.thread.is_alive()
-            raise TimeoutError(
-                f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. "
-                f"Read thread alive: {thread_alive}."
-            )
+        if self.latest_frame is None:
+            self.latest_frame = self.frame_queue.get()
+            self.frame_queue.task_done()
+            return self.latest_frame
 
-        with self.frame_lock:
-            frame = self.latest_frame
-            self.new_frame_event.clear()
+        try:
+            frame = self.frame_queue.get(timeout=timeout_ms / 1000.0)
+            self.frame_queue.task_done()
+        except Empty:
+            process_alive = self.process is not None and self.process.is_alive()
+            if process_alive:
+                logger.warning(f"{self} async_read timed out after {timeout_ms} ms but camera is still running.")
+                return self.latest_frame
+            else:
+                raise TimeoutError(
+                    f"{self} async_read timed out after {timeout_ms} ms: camera is not responding !"
+                )
 
         if frame is None:
             raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
-
-        return frame
+        else:
+            self.latest_frame = frame
+            return self.latest_frame
 
     def disconnect(self) -> None:
         """

src/lerobot/datasets/dataset_tools.py

@@ -39,6 +39,7 @@ from lerobot.datasets.aggregate import aggregate_datasets
 from lerobot.datasets.compute_stats import aggregate_stats
 from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
 from lerobot.datasets.utils import (
+    DATA_DIR,
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_DATA_PATH,
@@ -962,28 +963,23 @@ def _copy_data_with_feature_changes(
     remove_features: list[str] | None = None,
 ) -> None:
     """Copy data while adding or removing features."""
-    if dataset.meta.episodes is None:
-        dataset.meta.episodes = load_episodes(dataset.meta.root)
+    data_dir = dataset.root / DATA_DIR
+    parquet_files = sorted(data_dir.glob("*/*.parquet"))
 
-    # Map file paths to episode indices to extract chunk/file indices
-    file_to_episodes: dict[Path, set[int]] = {}
-    for ep_idx in range(dataset.meta.total_episodes):
-        file_path = dataset.meta.get_data_file_path(ep_idx)
-        if file_path not in file_to_episodes:
-            file_to_episodes[file_path] = set()
-        file_to_episodes[file_path].add(ep_idx)
+    if not parquet_files:
+        raise ValueError(f"No parquet files found in {data_dir}")
 
     frame_idx = 0
 
-    for src_path in tqdm(sorted(file_to_episodes.keys()), desc="Processing data files"):
-        df = pd.read_parquet(dataset.root / src_path).reset_index(drop=True)
+    for src_path in tqdm(parquet_files, desc="Processing data files"):
+        df = pd.read_parquet(src_path).reset_index(drop=True)
 
-        # Get chunk_idx and file_idx from the source file's first episode
-        episodes_in_file = file_to_episodes[src_path]
-        first_ep_idx = min(episodes_in_file)
-        src_ep = dataset.meta.episodes[first_ep_idx]
-        chunk_idx = src_ep["data/chunk_index"]
-        file_idx = src_ep["data/file_index"]
+        relative_path = src_path.relative_to(dataset.root)
+        chunk_dir = relative_path.parts[1]
+        file_name = relative_path.parts[2]
+
+        chunk_idx = int(chunk_dir.split("-")[1])
+        file_idx = int(file_name.split("-")[1].split(".")[0])
 
         if remove_features:
             df = df.drop(columns=remove_features, errors="ignore")
@@ -1009,7 +1005,7 @@ def _copy_data_with_feature_changes(
                         df[feature_name] = feature_slice
             frame_idx = end_idx
 
-        # Write using the preserved chunk_idx and file_idx from source
+        # Write using the same chunk/file structure as source
         dst_path = new_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         dst_path.parent.mkdir(parents=True, exist_ok=True)