linux only

lerobot/common/robot_devices/cameras/opencv.py

@@ -50,6 +50,24 @@ def find_camera_indices(raise_when_empty=False, max_index_search_range=MAX_OPENC
 
     return camera_ids
 
+def find_camera_ports(raise_when_empty=False):
+    camera_ports = []
+    for port in Path("/dev").glob("video*"):
+        port = str(port)
+        camera = cv2.VideoCapture(port)
+        is_open = camera.isOpened()
+        camera.release()
+
+        if is_open:
+            print(f"Camera found at port {port}")
+            camera_ports.append(port)
+
+    if raise_when_empty and len(camera_ports) == 0:
+        raise OSError(
+            "Not a single camera was detected. Try re-plugging, or re-installing `opencv2`, or your camera driver, or make sure your camera is compatible with opencv2."
+        )
+
+    return camera_ports
 
 def save_image(img_array, camera_index, frame_index, images_dir):
     img = Image.fromarray(img_array)
@@ -59,11 +77,16 @@ def save_image(img_array, camera_index, frame_index, images_dir):
 
 
 def save_images_from_cameras(
-    images_dir: Path, camera_ids=None, fps=None, width=None, height=None, record_time_s=2
+    images_dir: Path, camera_ids: list[int] | list[str] | None=None, fps=None, width=None, height=None, record_time_s=2
 ):
+    
     if camera_ids is None:
-        print("Finding available camera indices")
-        camera_ids = find_camera_indices()
+        print("Finding available camera indices through '/dev/video*' ports")
+        camera_ids = find_camera_ports()
+
+        if len(camera_ids) == 0:
+            print(f"Finding available camera indices through scanning all indices from 0 to {MAX_OPENCV_INDEX}")
+            camera_ids = find_camera_indices()
 
     print("Connecting cameras")
     cameras = []
@@ -76,7 +99,7 @@ def save_images_from_cameras(
         cameras.append(camera)
 
     images_dir = Path(
-        images_dir,
+        images_dir
     )
     if images_dir.exists():
         shutil.rmtree(
@@ -99,7 +122,7 @@ def save_images_from_cameras(
                 executor.submit(
                     save_image,
                     image,
-                    camera.camera_index,
+                    camera.camera_index_to_int(),
                     frame_index,
                     images_dir,
                 )
@@ -182,7 +205,7 @@ class OpenCVCamera:
     ```
     """
 
-    def __init__(self, camera_index: int, config: OpenCVCameraConfig | None = None, **kwargs):
+    def __init__(self, camera_index: int | str, config: OpenCVCameraConfig | None = None, **kwargs):
         if config is None:
             config = OpenCVCameraConfig()
         # Overwrite config arguments using kwargs
@@ -194,17 +217,17 @@ class OpenCVCamera:
         self.height = config.height
         self.color_mode = config.color_mode
 
-        if not isinstance(self.camera_index, int):
-            raise ValueError(
-                f"Camera index must be provided as an int, but {self.camera_index} was given instead."
-            )
-
         self.camera = None
         self.is_connected = False
         self.thread = None
         self.stop_event = None
         self.color_image = None
         self.logs = {}
+    
+    def camera_index_to_int(self):
+        if isinstance(self.camera_index, int):
+            return self.camera_index
+        return int(self.camera_index.replace("/dev/video", ""))
 
     def connect(self):
         if self.is_connected:
@@ -221,7 +244,9 @@ class OpenCVCamera:
         # valid cameras.
         if not is_camera_open:
             # Verify that the provided `camera_index` is valid before printing the traceback
-            available_cam_ids = find_camera_indices()
+            available_cam_ids = find_camera_ports()
+            if len(available_cam_ids) == 0:
+                available_cam_ids = find_camera_indices()
             if self.camera_index not in available_cam_ids:
                 raise ValueError(
                     f"`camera_index` is expected to be one of these available cameras {available_cam_ids}, but {self.camera_index} is provided instead."
@@ -365,7 +390,7 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--camera-ids",
-        type=int,
+        type=str,
         nargs="*",
         default=None,
         help="List of camera indices used to instantiate the `OpenCVCamera`. If not provided, find and use all available camera indices.",

lerobot/common/robot_devices/cameras/utils.py

@@ -2,6 +2,7 @@ from pathlib import Path
 from typing import Protocol
 
 import cv2
+import einops
 import numpy as np
 
 
@@ -39,6 +40,16 @@ def save_depth_image(depth, path, write_shape=False):
     cv2.imwrite(str(path), depth_image)
 
 
+def convert_torch_image_to_cv2(tensor, rgb_to_bgr=True):
+    assert tensor.ndim == 3
+    c, h, w = tensor.shape
+    assert c < h and c < w
+    color_image = einops.rearrange(tensor, "c h w -> h w c").numpy()
+    if rgb_to_bgr:
+        color_image = cv2.cvtColor(color_image, cv2.COLOR_RGB2BGR)
+    return color_image
+
+
 # Defines a camera type
 class Camera(Protocol):
     def connect(self): ...

lerobot/common/robot_devices/robots/factory.py

@@ -9,7 +9,7 @@ def make_robot(name):
         robot = KochRobot(
             leader_arms={
                 "main": DynamixelMotorsBus(
-                    port="/dev/tty.usbmodem575E0031751",
+                    port="/dev/ttyACM1",
                     motors={
                         # name: (index, model)
                         "shoulder_pan": (1, "xl330-m077"),
@@ -23,7 +23,7 @@ def make_robot(name):
             },
             follower_arms={
                 "main": DynamixelMotorsBus(
-                    port="/dev/tty.usbmodem575E0032081",
+                    port="/dev/ttyACM0",
                     motors={
                         # name: (index, model)
                         "shoulder_pan": (1, "xl430-w250"),
@@ -36,8 +36,8 @@ def make_robot(name):
                 ),
             },
             cameras={
-                "laptop": OpenCVCamera(0, fps=30, width=640, height=480),
-                "phone": OpenCVCamera(1, fps=30, width=640, height=480),
+                "laptop": OpenCVCamera("/dev/video2", fps=30, width=640, height=480),
+                "phone": OpenCVCamera("/dev/video6", fps=30, width=640, height=480),
             },
         )
     else:

lerobot/configs/default.yaml

@@ -72,13 +72,13 @@ training:
       min_max: [0.8, 1.2]
     saturation:
       weight: 1
-      min_max: [0.5, 1.5]
+      min_max: [0.7, 1.3]
     hue:
       weight: 1
-      min_max: [-0.05, 0.05]
+      min_max: [-0.02, 0.02]
     sharpness:
       weight: 1
-      min_max: [0.8, 1.2]
+      min_max: [0.1, 1.9]
 
 eval:
   n_episodes: 1

lerobot/scripts/control_robot.py

@@ -88,8 +88,10 @@ import platform
 import shutil
 import time
 from contextlib import nullcontext
+from functools import cache
 from pathlib import Path
 
+import cv2
 import torch
 import tqdm
 from huggingface_hub import create_branch
@@ -105,6 +107,7 @@ from lerobot.common.datasets.push_dataset_to_hub.utils import concatenate_episod
 from lerobot.common.datasets.utils import calculate_episode_data_index
 from lerobot.common.datasets.video_utils import encode_video_frames
 from lerobot.common.policies.factory import make_policy
+from lerobot.common.robot_devices.cameras.utils import convert_torch_image_to_cv2
 from lerobot.common.robot_devices.robots.factory import make_robot
 from lerobot.common.robot_devices.robots.utils import Robot
 from lerobot.common.utils.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
@@ -179,7 +182,8 @@ def log_control_info(robot, dt_s, episode_index=None, frame_index=None, fps=None
     logging.info(info_str)
 
 
-def get_is_headless():
+@cache
+def is_headless():
     if platform.system() == "Linux":
         display = os.environ.get("DISPLAY")
         if display is None or display == "":
@@ -255,32 +259,10 @@ def record_dataset(
     else:
         episode_index = 0
 
-    is_headless = get_is_headless()
-
-    # Execute a few seconds without recording data, to give times
-    # to the robot devices to connect and start synchronizing.
-    timestamp = 0
-    start_time = time.perf_counter()
-    is_warmup_print = False
-    while timestamp < warmup_time_s:
-        if not is_warmup_print:
-            logging.info("Warming up (no data recording)")
-            os.system('say "Warmup" &')
-            is_warmup_print = True
-
-        now = time.perf_counter()
-        observation, action = robot.teleop_step(record_data=True)
-
-        if not is_headless:
-            image_keys = [key for key in observation if "image" in key]
-
-        dt_s = time.perf_counter() - now
-        busy_wait(1 / fps - dt_s)
-
-        dt_s = time.perf_counter() - now
-        log_control_info(robot, dt_s, fps=fps)
-
-        timestamp = time.perf_counter() - start_time
+    if is_headless():
+        logging.info(
+            "Headless environment detected. Display cameras on screen and keyboard inputs will not be available."
+        )
 
     # Allow to exit early while recording an episode or resetting the environment,
     # by tapping the right arrow key '->'. This might require a sudo permission
@@ -290,9 +272,7 @@ def record_dataset(
     stop_recording = False
 
     # Only import pynput if not in a headless environment
-    if is_headless:
-        logging.info("Headless environment detected. Keyboard input will not be available.")
-    else:
+    if not is_headless():
         from pynput import keyboard
 
         def on_press(key):
@@ -315,6 +295,35 @@ def record_dataset(
         listener = keyboard.Listener(on_press=on_press)
         listener.start()
 
+    # Execute a few seconds without recording data, to give times
+    # to the robot devices to connect and start synchronizing.
+    timestamp = 0
+    start_time = time.perf_counter()
+    is_warmup_print = False
+    while timestamp < warmup_time_s:
+        if not is_warmup_print:
+            logging.info("Warming up (no data recording)")
+            os.system('say "Warmup" &')
+            is_warmup_print = True
+
+        now = time.perf_counter()
+
+        observation, action = robot.teleop_step(record_data=True)
+
+        if not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+            cv2.waitKey(1)
+
+        dt_s = time.perf_counter() - now
+        busy_wait(1 / fps - dt_s)
+
+        dt_s = time.perf_counter() - now
+        log_control_info(robot, dt_s, fps=fps)
+
+        timestamp = time.perf_counter() - start_time
+
     # Save images using threads to reach high fps (30 and more)
     # Using `with` to exist smoothly if an execption is raised.
     # Using only 4 worker threads to avoid blocking the main thread.
@@ -342,6 +351,12 @@ def record_dataset(
                         )
                     ]
 
+                if not is_headless():
+                    image_keys = [key for key in observation if "image" in key]
+                    for key in image_keys:
+                        cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+                    cv2.waitKey(1)
+
                 for key in not_image_keys:
                     if key not in ep_dict:
                         ep_dict[key] = []
@@ -434,7 +449,7 @@ def record_dataset(
             if is_last_episode:
                 logging.info("Done recording")
                 os.system('say "Done recording"')
-                if not is_headless:
+                if not is_headless():
                     listener.stop()
 
                 logging.info("Waiting for threads writing the images on disk to terminate...")
@@ -444,6 +459,10 @@ def record_dataset(
                     pass
                 break
 
+    robot.disconnect()
+    if not is_headless():
+        cv2.destroyAllWindows()
+
     num_episodes = episode_index
 
     logging.info("Encoding videos")
@@ -493,6 +512,7 @@ def record_dataset(
         stats = compute_stats(lerobot_dataset)
         lerobot_dataset.stats = stats
     else:
+        stats = {}
         logging.info("Skipping computation of the dataset statistrics")
 
     hf_dataset = hf_dataset.with_format(None)  # to remove transforms that cant be saved
@@ -544,7 +564,14 @@ def replay_episode(robot: Robot, episode: int, fps: int | None = None, root="dat
         log_control_info(robot, dt_s, fps=fps)
 
 
-def run_policy(robot: Robot, policy: torch.nn.Module, hydra_cfg: DictConfig, run_time_s: float | None = None):
+def run_policy(
+    robot: Robot,
+    policy: torch.nn.Module,
+    hydra_cfg: DictConfig,
+    warmup_time_s: float = 4,
+    run_time_s: float | None = None,
+    reset_time_s: float = 15,
+):
     # TODO(rcadene): Add option to record eval dataset and logs
 
     # Check device is available
@@ -562,12 +589,76 @@ def run_policy(robot: Robot, policy: torch.nn.Module, hydra_cfg: DictConfig, run
     if not robot.is_connected:
         robot.connect()
 
+    if is_headless():
+        logging.info(
+            "Headless environment detected. Display cameras on screen and keyboard inputs will not be available."
+        )
+
+    # Allow to reset environment or exit early
+    # by tapping the right arrow key '->'. This might require a sudo permission
+    # to allow your terminal to monitor keyboard events.
+    reset_environment = False
+    exit_reset = False
+
+    # Only import pynput if not in a headless environment
+    if not is_headless():
+        from pynput import keyboard
+
+        def on_press(key):
+            nonlocal reset_environment, exit_reset
+            try:
+                if key == keyboard.Key.right and not reset_environment:
+                    print("Right arrow key pressed. Suspend robot control to reset environment...")
+                    reset_environment = True
+                elif key == keyboard.Key.right and reset_environment:
+                    print("Right arrow key pressed. Enable robot control and exit reset environment...")
+                    exit_reset = True
+            except Exception as e:
+                print(f"Error handling key press: {e}")
+
+        listener = keyboard.Listener(on_press=on_press)
+        listener.start()
+
+    # Execute a few seconds without recording data, to give times
+    # to the robot devices to connect and start synchronizing.
+    timestamp = 0
+    start_time = time.perf_counter()
+    is_warmup_print = False
+    while timestamp < warmup_time_s:
+        if not is_warmup_print:
+            logging.info("Warming up (no data recording)")
+            os.system('say "Warmup" &')
+            is_warmup_print = True
+
+        now = time.perf_counter()
+        observation = robot.capture_observation()
+
+        if not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, convert_torch_image_to_cv2(observation[key]))
+            cv2.waitKey(1)
+
+        dt_s = time.perf_counter() - now
+        busy_wait(1 / fps - dt_s)
+
+        dt_s = time.perf_counter() - now
+        log_control_info(robot, dt_s, fps=fps)
+
+        timestamp = time.perf_counter() - start_time
+
     start_time = time.perf_counter()
     while True:
         now = time.perf_counter()
 
         observation = robot.capture_observation()
 
+        if not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(key, convert_torch_image_to_cv2(observation[key]))
+            cv2.waitKey(1)
+
         with (
             torch.inference_mode(),
             torch.autocast(device_type=device.type)
@@ -598,6 +689,25 @@ def run_policy(robot: Robot, policy: torch.nn.Module, hydra_cfg: DictConfig, run
         if run_time_s is not None and time.perf_counter() - start_time > run_time_s:
             break
 
+        if reset_environment:
+            # Start resetting env while the executor are finishing
+            logging.info("Reset the environment")
+            os.system('say "Reset the environment" &')
+
+            # Wait if necessary
+            timestamp = 0
+            start_time = time.perf_counter()
+            with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar:
+                while timestamp < reset_time_s:
+                    time.sleep(1)
+                    timestamp = time.perf_counter() - start_time
+                    pbar.update(1)
+                    if exit_reset:
+                        exit_reset = False
+                        break
+
+            reset_environment = False
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

lerobot/scripts/visualize_image_transforms.py

@@ -172,4 +172,4 @@ def visualize_transforms_cli(cfg):
 
 
 if __name__ == "__main__":
-    visualize_transforms()
+    visualize_transforms_cli()