Compare commits

...

2 Commits

Author SHA1 Message Date
CarolinePascal ee50b0f24b feat(WIP): adding support for multi-process camera video capture 2025-11-04 19:51:58 +01:00
Michel Aractingi f6b16f6d97 fix(dataset_tools) Critical bug in modify features (#2342)
* 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>
2025-11-04 15:56:41 +01:00
2 changed files with 52 additions and 48 deletions
+38 -30
View File
@@ -23,6 +23,8 @@ import platform
import time import time
from pathlib import Path from pathlib import Path
from threading import Event, Lock, Thread 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 typing import Any
from numpy.typing import NDArray # type: ignore # TODO: add type stubs for numpy.typing 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.videocapture: cv2.VideoCapture | None = None
self.thread: Thread | None = None self.process: Process | None = None
self.stop_event: Event | None = None self.stop_event: EventProcess | None = None
self.frame_lock: Lock = Lock() self.frame_queue: Queue = Queue()
self.latest_frame: NDArray[Any] | None = None self.latest_frame: NDArray[Any] | None = None
self.new_frame_event: Event = Event()
self.rotation: int | None = get_cv2_rotation(config.rotation) self.rotation: int | None = get_cv2_rotation(config.rotation)
self.backend: int = get_cv2_backend() self.backend: int = get_cv2_backend()
@@ -442,37 +443,36 @@ class OpenCVCamera(Camera):
while not self.stop_event.is_set(): while not self.stop_event.is_set():
try: try:
color_image = self.read() color_image = self.read()
self.frame_queue.put_nowait(color_image)
with self.frame_lock:
self.latest_frame = color_image
self.new_frame_event.set()
except DeviceNotConnectedError: except DeviceNotConnectedError:
break break
except Exception as e: except Exception as e:
logger.warning(f"Error reading frame in background thread for {self}: {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.""" """Starts or restarts the background read thread if it's not running."""
if self.thread is not None and self.thread.is_alive(): if self.process is not None and self.process.is_alive():
self.thread.join(timeout=0.1) self.frame_queue.join()
self.process.join()
if self.stop_event is not None: if self.stop_event is not None:
self.stop_event.set() self.stop_event.set()
self.stop_event = Event() self.stop_event = Event()
self.thread = Thread(target=self._read_loop, args=(), name=f"{self}_read_loop") self.process = Process(target=self._read_loop, args=(), name=f"{self}_read_loop")
self.thread.daemon = True self.process.daemon = True
self.thread.start() self.process.start()
def _stop_read_thread(self) -> None: def _stop_read_thread(self) -> None:
"""Signals the background read thread to stop and waits for it to join.""" """Signals the background read thread to stop and waits for it to join."""
if self.stop_event is not None: if self.stop_event is not None:
self.stop_event.set() self.stop_event.set()
if self.thread is not None and self.thread.is_alive(): if self.process is not None and self.process.is_alive():
self.thread.join(timeout=2.0) self.frame_queue.join()
self.process.join()
self.thread = None self.process = None
self.stop_event = None self.stop_event = None
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]: def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
@@ -499,24 +499,32 @@ class OpenCVCamera(Camera):
if not self.is_connected: if not self.is_connected:
raise DeviceNotConnectedError(f"{self} is not connected.") raise DeviceNotConnectedError(f"{self} is not connected.")
if self.thread is None or not self.thread.is_alive(): if self.process is None or not self.process.is_alive():
self._start_read_thread() self._start_read_process()
if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): if self.latest_frame is None:
thread_alive = self.thread is not None and self.thread.is_alive() self.latest_frame = self.frame_queue.get()
raise TimeoutError( self.frame_queue.task_done()
f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " return self.latest_frame
f"Read thread alive: {thread_alive}."
)
with self.frame_lock: try:
frame = self.latest_frame frame = self.frame_queue.get(timeout=timeout_ms / 1000.0)
self.new_frame_event.clear() 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: if frame is None:
raise RuntimeError(f"Internal error: Event set but no frame available for {self}.") raise RuntimeError(f"Internal error: Event set but no frame available for {self}.")
else:
return frame self.latest_frame = frame
return self.latest_frame
def disconnect(self) -> None: def disconnect(self) -> None:
""" """
+14 -18
View File
@@ -39,6 +39,7 @@ from lerobot.datasets.aggregate import aggregate_datasets
from lerobot.datasets.compute_stats import aggregate_stats from lerobot.datasets.compute_stats import aggregate_stats
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
from lerobot.datasets.utils import ( from lerobot.datasets.utils import (
DATA_DIR,
DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_SIZE,
DEFAULT_DATA_FILE_SIZE_IN_MB, DEFAULT_DATA_FILE_SIZE_IN_MB,
DEFAULT_DATA_PATH, DEFAULT_DATA_PATH,
@@ -962,28 +963,23 @@ def _copy_data_with_feature_changes(
remove_features: list[str] | None = None, remove_features: list[str] | None = None,
) -> None: ) -> None:
"""Copy data while adding or removing features.""" """Copy data while adding or removing features."""
if dataset.meta.episodes is None: data_dir = dataset.root / DATA_DIR
dataset.meta.episodes = load_episodes(dataset.meta.root) parquet_files = sorted(data_dir.glob("*/*.parquet"))
# Map file paths to episode indices to extract chunk/file indices if not parquet_files:
file_to_episodes: dict[Path, set[int]] = {} raise ValueError(f"No parquet files found in {data_dir}")
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)
frame_idx = 0 frame_idx = 0
for src_path in tqdm(sorted(file_to_episodes.keys()), desc="Processing data files"): for src_path in tqdm(parquet_files, desc="Processing data files"):
df = pd.read_parquet(dataset.root / src_path).reset_index(drop=True) df = pd.read_parquet(src_path).reset_index(drop=True)
# Get chunk_idx and file_idx from the source file's first episode relative_path = src_path.relative_to(dataset.root)
episodes_in_file = file_to_episodes[src_path] chunk_dir = relative_path.parts[1]
first_ep_idx = min(episodes_in_file) file_name = relative_path.parts[2]
src_ep = dataset.meta.episodes[first_ep_idx]
chunk_idx = src_ep["data/chunk_index"] chunk_idx = int(chunk_dir.split("-")[1])
file_idx = src_ep["data/file_index"] file_idx = int(file_name.split("-")[1].split(".")[0])
if remove_features: if remove_features:
df = df.drop(columns=remove_features, errors="ignore") df = df.drop(columns=remove_features, errors="ignore")
@@ -1009,7 +1005,7 @@ def _copy_data_with_feature_changes(
df[feature_name] = feature_slice df[feature_name] = feature_slice
frame_idx = end_idx 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 = new_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
dst_path.parent.mkdir(parents=True, exist_ok=True) dst_path.parent.mkdir(parents=True, exist_ok=True)