mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-12 04:21:45 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| ee50b0f24b | |||
| f6b16f6d97 |
@@ -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:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user