From fd8032a023e6e9d1e50638dac95061da2bb01199 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Mar 2025 12:12:42 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- lerobot/common/robot_devices/cameras/opencv.py | 16 ++++++++-------- lerobot/common/robot_devices/control_utils.py | 14 +++++++++++--- lerobot/common/robot_devices/robots/configs.py | 7 +++++-- lerobot/common/robot_devices/robots/roarm_m3.py | 13 ++++++------- lerobot/common/robot_devices/robots/utils.py | 3 ++- 5 files changed, 32 insertions(+), 21 deletions(-) diff --git a/lerobot/common/robot_devices/cameras/opencv.py b/lerobot/common/robot_devices/cameras/opencv.py index 506f6075..5a4e7feb 100644 --- a/lerobot/common/robot_devices/cameras/opencv.py +++ b/lerobot/common/robot_devices/cameras/opencv.py @@ -31,20 +31,20 @@ from lerobot.common.utils.utils import capture_timestamp_utc MAX_OPENCV_INDEX = 60 undistort = True + + def undistort_image(image): import cv2 - camera_matrix = np.array([ - [289.11451, 0., 347.23664], - [0., 289.75319, 235.67429], - [0., 0., 1.] - ]) - + + camera_matrix = np.array([[289.11451, 0.0, 347.23664], [0.0, 289.75319, 235.67429], [0.0, 0.0, 1.0]]) + dist_coeffs = np.array([-0.208848, 0.028006, -0.000705, -0.000820, 0.000000]) - + undistorted_image = cv2.undistort(image, camera_matrix, dist_coeffs) - + return undistorted_image + def find_cameras(raise_when_empty=False, max_index_search_range=MAX_OPENCV_INDEX, mock=False) -> list[dict]: cameras = [] if platform.system() == "Linux": diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index c5870e1f..dfc4b1cb 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -153,13 +153,16 @@ def predict_action(observation, policy, device, use_amp): # return listener, events + def init_keyboard_listener(): events = {} events["exit_early"] = False events["rerecord_episode"] = False events["stop_recording"] = False + import threading + from sshkeyboard import listen_keyboard - import threading + def on_press(key): try: if key == "right": @@ -179,7 +182,8 @@ def init_keyboard_listener(): listener = threading.Thread(target=listen_keyboard, kwargs={"on_press": on_press}) listener.start() - return listener,events + return listener, events + def warmup_record( robot, @@ -284,7 +288,7 @@ def control_loop( dataset.add_frame(frame) if display_cameras: - # if display_cameras and not is_headless(): + # if display_cameras and 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)) @@ -316,6 +320,7 @@ def reset_environment(robot, events, reset_time_s, fps): teleoperate=True, ) + # def stop_recording(robot, listener, display_cameras): # robot.disconnect() @@ -326,16 +331,19 @@ def reset_environment(robot, events, reset_time_s, fps): # if display_cameras: # cv2.destroyAllWindows() + def stop_recording(robot, listener, display_cameras): robot.disconnect() from sshkeyboard import stop_listening + if listener is not None: stop_listening() if display_cameras: cv2.destroyAllWindows() + def sanity_check_dataset_name(repo_id, policy_cfg): _, dataset_name = repo_id.split("/") # either repo_id doesnt start with "eval_" and there is no policy diff --git a/lerobot/common/robot_devices/robots/configs.py b/lerobot/common/robot_devices/robots/configs.py index ac38c7eb..f66a9bef 100644 --- a/lerobot/common/robot_devices/robots/configs.py +++ b/lerobot/common/robot_devices/robots/configs.py @@ -479,6 +479,7 @@ class So100RobotConfig(ManipulatorRobotConfig): mock: bool = False + @RobotConfig.register_subclass("roarm_m3") @dataclass class RoarmRobotConfig(RobotConfig): @@ -489,14 +490,14 @@ class RoarmRobotConfig(RobotConfig): leader_arms: dict[str, str] = field( default_factory=lambda: { - "main": "/dev/ttyUSB0", + "main": "/dev/ttyUSB0", } ) # Follower arms configuration: left and right ports follower_arms: dict[str, str] = field( default_factory=lambda: { - "main": "/dev/ttyUSB1", + "main": "/dev/ttyUSB1", } ) @@ -519,6 +520,7 @@ class RoarmRobotConfig(RobotConfig): mock: bool = False + @RobotConfig.register_subclass("stretch") @dataclass class StretchRobotConfig(RobotConfig): @@ -554,6 +556,7 @@ class StretchRobotConfig(RobotConfig): mock: bool = False + @RobotConfig.register_subclass("lekiwi") @dataclass class LeKiwiRobotConfig(RobotConfig): diff --git a/lerobot/common/robot_devices/robots/roarm_m3.py b/lerobot/common/robot_devices/robots/roarm_m3.py index 24461194..f6b8e468 100644 --- a/lerobot/common/robot_devices/robots/roarm_m3.py +++ b/lerobot/common/robot_devices/robots/roarm_m3.py @@ -4,19 +4,17 @@ and send orders to its motors. # TODO(rcadene, aliberts): reorganize the codebase into one file per robot, with the associated # calibration procedure, to make it easy for people to add their own robot. -import json import logging import time -import warnings -from pathlib import Path import numpy as np import torch +from roarm_sdk.roarm import roarm from lerobot.common.robot_devices.cameras.utils import make_cameras_from_configs from lerobot.common.robot_devices.robots.configs import RoarmRobotConfig from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError -from roarm_sdk.roarm import roarm + def ensure_safe_goal_position( goal_pos: torch.Tensor, present_pos: torch.Tensor, max_relative_target: float | list[float] @@ -37,7 +35,8 @@ def ensure_safe_goal_position( return safe_goal_pos -def make_roarm_from_configs(configs: dict[str, str]) -> (dict[str, roarm]): + +def make_roarm_from_configs(configs: dict[str, str]) -> dict[str, roarm]: roarms = {} for key, port in configs.items(): @@ -45,8 +44,8 @@ def make_roarm_from_configs(configs: dict[str, str]) -> (dict[str, roarm]): return roarms -class RoarmRobot: +class RoarmRobot: def __init__( self, config: RoarmRobotConfig, @@ -58,7 +57,7 @@ class RoarmRobot: self.cameras = make_cameras_from_configs(self.config.cameras) self.is_connected = False self.logs = {} - + @property def camera_features(self) -> dict: cam_ft = {} diff --git a/lerobot/common/robot_devices/robots/utils.py b/lerobot/common/robot_devices/robots/utils.py index d418ec34..07714cff 100644 --- a/lerobot/common/robot_devices/robots/utils.py +++ b/lerobot/common/robot_devices/robots/utils.py @@ -7,9 +7,9 @@ from lerobot.common.robot_devices.robots.configs import ( LeKiwiRobotConfig, ManipulatorRobotConfig, MossRobotConfig, + RoarmRobotConfig, RobotConfig, So100RobotConfig, - RoarmRobotConfig, StretchRobotConfig, ) @@ -73,6 +73,7 @@ def make_robot_from_config(config: RobotConfig): return StretchRobot(config) + def make_robot(robot_type: str, **kwargs) -> Robot: config = make_robot_config(robot_type, **kwargs) return make_robot_from_config(config)