fix: camera and motors modules for mock robots

lerobot/common/mocks/__init__.py

@@ -0,0 +1 @@
+# Common mocks for robot devices and testing

lerobot/common/mocks/cameras/mock_cv2.py

@@ -0,0 +1,101 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import cache
+
+import numpy as np
+
+CAP_V4L2 = 200
+CAP_DSHOW = 700
+CAP_AVFOUNDATION = 1200
+CAP_ANY = -1
+
+CAP_PROP_FPS = 5
+CAP_PROP_FRAME_WIDTH = 3
+CAP_PROP_FRAME_HEIGHT = 4
+COLOR_RGB2BGR = 4
+COLOR_BGR2RGB = 4
+
+ROTATE_90_COUNTERCLOCKWISE = 2
+ROTATE_90_CLOCKWISE = 0
+ROTATE_180 = 1
+
+
+@cache
+def _generate_image(width: int, height: int):
+    return np.random.randint(0, 256, size=(height, width, 3), dtype=np.uint8)
+
+
+def cvtColor(color_image, color_conversion):  # noqa: N802
+    if color_conversion in [COLOR_RGB2BGR, COLOR_BGR2RGB]:
+        return color_image[:, :, [2, 1, 0]]
+    else:
+        raise NotImplementedError(color_conversion)
+
+
+def rotate(color_image, rotation):
+    if rotation is None:
+        return color_image
+    elif rotation == ROTATE_90_CLOCKWISE:
+        return np.rot90(color_image, k=1)
+    elif rotation == ROTATE_180:
+        return np.rot90(color_image, k=2)
+    elif rotation == ROTATE_90_COUNTERCLOCKWISE:
+        return np.rot90(color_image, k=3)
+    else:
+        raise NotImplementedError(rotation)
+
+
+class VideoCapture:
+    def __init__(self, *args, **kwargs):
+        self._mock_dict = {
+            CAP_PROP_FPS: 30,
+            CAP_PROP_FRAME_WIDTH: 640,
+            CAP_PROP_FRAME_HEIGHT: 480,
+        }
+        self._is_opened = True
+
+    def isOpened(self):  # noqa: N802
+        return self._is_opened
+
+    def set(self, propId: int, value: float) -> bool:  # noqa: N803
+        if not self._is_opened:
+            raise RuntimeError("Camera is not opened")
+        self._mock_dict[propId] = value
+        return True
+
+    def get(self, propId: int) -> float:  # noqa: N803
+        if not self._is_opened:
+            raise RuntimeError("Camera is not opened")
+        value = self._mock_dict[propId]
+        if value == 0:
+            if propId == CAP_PROP_FRAME_HEIGHT:
+                value = 480
+            elif propId == CAP_PROP_FRAME_WIDTH:
+                value = 640
+        return value
+
+    def read(self):
+        if not self._is_opened:
+            raise RuntimeError("Camera is not opened")
+        h = self.get(CAP_PROP_FRAME_HEIGHT)
+        w = self.get(CAP_PROP_FRAME_WIDTH)
+        ret = True
+        return ret, _generate_image(width=w, height=h)
+
+    def release(self):
+        self._is_opened = False
+
+    def __del__(self):
+        if self._is_opened:
+            self.release()

lerobot/common/mocks/cameras/mock_pyrealsense2.py

@@ -0,0 +1,148 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import enum
+
+import numpy as np
+
+
+class stream(enum.Enum):  # noqa: N801
+    color = 0
+    depth = 1
+
+
+class format(enum.Enum):  # noqa: N801
+    rgb8 = 0
+    z16 = 1
+
+
+class config:  # noqa: N801
+    def enable_device(self, device_id: str):
+        self.device_enabled = device_id
+
+    def enable_stream(self, stream_type: stream, width=None, height=None, color_format=None, fps=None):
+        self.stream_type = stream_type
+        # Overwrite default values when possible
+        self.width = 848 if width is None else width
+        self.height = 480 if height is None else height
+        self.color_format = format.rgb8 if color_format is None else color_format
+        self.fps = 30 if fps is None else fps
+
+
+class RSColorProfile:
+    def __init__(self, config):
+        self.config = config
+
+    def fps(self):
+        return self.config.fps
+
+    def width(self):
+        return self.config.width
+
+    def height(self):
+        return self.config.height
+
+
+class RSColorStream:
+    def __init__(self, config):
+        self.config = config
+
+    def as_video_stream_profile(self):
+        return RSColorProfile(self.config)
+
+
+class RSProfile:
+    def __init__(self, config):
+        self.config = config
+
+    def get_stream(self, color_format):
+        del color_format  # unused
+        return RSColorStream(self.config)
+
+
+class pipeline:  # noqa: N801
+    def __init__(self):
+        self.started = False
+        self.config = None
+
+    def start(self, config):
+        self.started = True
+        self.config = config
+        return RSProfile(self.config)
+
+    def stop(self):
+        if not self.started:
+            raise RuntimeError("You need to start the camera before stop.")
+        self.started = False
+        self.config = None
+
+    def wait_for_frames(self, timeout_ms=50000):
+        del timeout_ms  # unused
+        return RSFrames(self.config)
+
+
+class RSFrames:
+    def __init__(self, config):
+        self.config = config
+
+    def get_color_frame(self):
+        return RSColorFrame(self.config)
+
+    def get_depth_frame(self):
+        return RSDepthFrame(self.config)
+
+
+class RSColorFrame:
+    def __init__(self, config):
+        self.config = config
+
+    def get_data(self):
+        data = np.ones((self.config.height, self.config.width, 3), dtype=np.uint8)
+        # Create a difference between rgb and bgr
+        data[:, :, 0] = 2
+        return data
+
+
+class RSDepthFrame:
+    def __init__(self, config):
+        self.config = config
+
+    def get_data(self):
+        return np.ones((self.config.height, self.config.width), dtype=np.uint16)
+
+
+class RSDevice:
+    def __init__(self):
+        pass
+
+    def get_info(self, camera_info) -> str:
+        del camera_info  # unused
+        # return fake serial number
+        return "123456789"
+
+
+class context:  # noqa: N801
+    def __init__(self):
+        pass
+
+    def query_devices(self):
+        return [RSDevice()]
+
+
+class camera_info:  # noqa: N801
+    # fake name
+    name = "Intel RealSense D435I"
+
+    def __init__(self, serial_number):
+        del serial_number
+        pass

lerobot/common/mocks/motors/__init__.py

@@ -0,0 +1 @@
+# Mocks for motor modules

lerobot/common/mocks/motors/mock_dynamixel_sdk.py

@@ -0,0 +1,107 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Mocked classes and functions from dynamixel_sdk to allow for continuous integration
+and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
+
+Warning: These mocked versions are minimalist. They do not exactly mock every behaviors
+from the original classes and functions (e.g. return types might be None instead of boolean).
+"""
+
+# from dynamixel_sdk import COMM_SUCCESS
+
+DEFAULT_BAUDRATE = 9_600
+COMM_SUCCESS = 0  # tx or rx packet communication success
+
+
+def convert_to_bytes(value, bytes):
+    # TODO(rcadene): remove need to mock `convert_to_bytes` by implemented the inverse transform
+    # `convert_bytes_to_value`
+    del bytes  # unused
+    return value
+
+
+def get_default_motor_values(motor_index):
+    return {
+        # Key (int) are from X_SERIES_CONTROL_TABLE
+        7: motor_index,  # ID
+        8: DEFAULT_BAUDRATE,  # Baud_rate
+        10: 0,  # Drive_Mode
+        64: 0,  # Torque_Enable
+        # Set 2560 since calibration values for Aloha gripper is between start_pos=2499 and end_pos=3144
+        # For other joints, 2560 will be autocorrected to be in calibration range
+        132: 2560,  # Present_Position
+    }
+
+
+class PortHandler:
+    def __init__(self, port):
+        self.port = port
+        # factory default baudrate
+        self.baudrate = DEFAULT_BAUDRATE
+
+    def openPort(self):  # noqa: N802
+        return True
+
+    def closePort(self):  # noqa: N802
+        pass
+
+    def setPacketTimeoutMillis(self, timeout_ms):  # noqa: N802
+        del timeout_ms  # unused
+
+    def getBaudRate(self):  # noqa: N802
+        return self.baudrate
+
+    def setBaudRate(self, baudrate):  # noqa: N802
+        self.baudrate = baudrate
+
+
+class PacketHandler:
+    def __init__(self, protocol_version):
+        del protocol_version  # unused
+        # Use packet_handler.data to communicate across Read and Write
+        self.data = {}
+
+
+class GroupSyncRead:
+    def __init__(self, port_handler, packet_handler, address, bytes):
+        self.packet_handler = packet_handler
+
+    def addParam(self, motor_index):  # noqa: N802
+        # Initialize motor default values
+        if motor_index not in self.packet_handler.data:
+            self.packet_handler.data[motor_index] = get_default_motor_values(motor_index)
+
+    def txRxPacket(self):  # noqa: N802
+        return COMM_SUCCESS
+
+    def getData(self, index, address, bytes):  # noqa: N802
+        return self.packet_handler.data[index][address]
+
+
+class GroupSyncWrite:
+    def __init__(self, port_handler, packet_handler, address, bytes):
+        self.packet_handler = packet_handler
+        self.address = address
+
+    def addParam(self, index, data):  # noqa: N802
+        # Initialize motor default values
+        if index not in self.packet_handler.data:
+            self.packet_handler.data[index] = get_default_motor_values(index)
+        self.changeParam(index, data)
+
+    def txPacket(self):  # noqa: N802
+        return COMM_SUCCESS
+
+    def changeParam(self, index, data):  # noqa: N802
+        self.packet_handler.data[index][self.address] = data

lerobot/common/mocks/motors/mock_scservo_sdk.py

@@ -0,0 +1,125 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Mocked classes and functions from dynamixel_sdk to allow for continuous integration
+and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
+
+Warning: These mocked versions are minimalist. They do not exactly mock every behaviors
+from the original classes and functions (e.g. return types might be None instead of boolean).
+"""
+
+# from dynamixel_sdk import COMM_SUCCESS
+
+DEFAULT_BAUDRATE = 1_000_000
+COMM_SUCCESS = 0  # tx or rx packet communication success
+
+
+def convert_to_bytes(value, bytes):
+    # TODO(rcadene): remove need to mock `convert_to_bytes` by implemented the inverse transform
+    # `convert_bytes_to_value`
+    del bytes  # unused
+    return value
+
+
+def get_default_motor_values(motor_index):
+    return {
+        # Key (int) are from SCS_SERIES_CONTROL_TABLE
+        5: motor_index,  # ID
+        6: DEFAULT_BAUDRATE,  # Baud_rate
+        10: 0,  # Drive_Mode
+        21: 32,  # P_Coefficient
+        22: 32,  # D_Coefficient
+        23: 0,  # I_Coefficient
+        40: 0,  # Torque_Enable
+        41: 254,  # Acceleration
+        31: -2047,  # Offset
+        33: 0,  # Mode
+        55: 1,  # Lock
+        # Set 2560 since calibration values for Aloha gripper is between start_pos=2499 and end_pos=3144
+        # For other joints, 2560 will be autocorrected to be in calibration range
+        56: 2560,  # Present_Position
+        58: 0,  # Present_Speed
+        69: 0,  # Present_Current
+        85: 150,  # Maximum_Acceleration
+    }
+
+
+class PortHandler:
+    def __init__(self, port):
+        self.port = port
+        # factory default baudrate
+        self.baudrate = DEFAULT_BAUDRATE
+        self.ser = SerialMock()
+
+    def openPort(self):  # noqa: N802
+        return True
+
+    def closePort(self):  # noqa: N802
+        pass
+
+    def setPacketTimeoutMillis(self, timeout_ms):  # noqa: N802
+        del timeout_ms  # unused
+
+    def getBaudRate(self):  # noqa: N802
+        return self.baudrate
+
+    def setBaudRate(self, baudrate):  # noqa: N802
+        self.baudrate = baudrate
+
+
+class PacketHandler:
+    def __init__(self, protocol_version):
+        del protocol_version  # unused
+        # Use packet_handler.data to communicate across Read and Write
+        self.data = {}
+
+
+class GroupSyncRead:
+    def __init__(self, port_handler, packet_handler, address, bytes):
+        self.packet_handler = packet_handler
+
+    def addParam(self, motor_index):  # noqa: N802
+        # Initialize motor default values
+        if motor_index not in self.packet_handler.data:
+            self.packet_handler.data[motor_index] = get_default_motor_values(motor_index)
+
+    def txRxPacket(self):  # noqa: N802
+        return COMM_SUCCESS
+
+    def getData(self, index, address, bytes):  # noqa: N802
+        return self.packet_handler.data[index][address]
+
+
+class GroupSyncWrite:
+    def __init__(self, port_handler, packet_handler, address, bytes):
+        self.packet_handler = packet_handler
+        self.address = address
+
+    def addParam(self, index, data):  # noqa: N802
+        if index not in self.packet_handler.data:
+            self.packet_handler.data[index] = get_default_motor_values(index)
+        self.changeParam(index, data)
+
+    def txPacket(self):  # noqa: N802
+        return COMM_SUCCESS
+
+    def changeParam(self, index, data):  # noqa: N802
+        self.packet_handler.data[index][self.address] = data
+
+
+class SerialMock:
+    def reset_output_buffer(self):
+        pass
+
+    def reset_input_buffer(self):
+        pass

lerobot/common/robot_devices/cameras/intelrealsense.py

@@ -48,7 +48,7 @@ def find_cameras(raise_when_empty=True, mock=False) -> list[dict]:
     connected to the computer.
     """
     if mock:
-        import tests.cameras.mock_pyrealsense2 as rs
+        import lerobot.common.mocks.cameras.mock_pyrealsense2 as rs
     else:
         import pyrealsense2 as rs
 
@@ -100,7 +100,7 @@ def save_images_from_cameras(
         serial_numbers = [cam["serial_number"] for cam in camera_infos]
 
     if mock:
-        import tests.cameras.mock_cv2 as cv2
+        import lerobot.common.mocks.cameras.mock_cv2 as cv2
     else:
         import cv2
 
@@ -253,7 +253,7 @@ class IntelRealSenseCamera:
         self.logs = {}
 
         if self.mock:
-            import tests.cameras.mock_cv2 as cv2
+            import lerobot.common.mocks.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -287,7 +287,7 @@ class IntelRealSenseCamera:
             )
 
         if self.mock:
-            import tests.cameras.mock_pyrealsense2 as rs
+            import lerobot.common.mocks.cameras.mock_pyrealsense2 as rs
         else:
             import pyrealsense2 as rs
 
@@ -375,7 +375,7 @@ class IntelRealSenseCamera:
             )
 
         if self.mock:
-            import tests.cameras.mock_cv2 as cv2
+            import lerobot.common.mocks.cameras.mock_cv2 as cv2
         else:
             import cv2
 

lerobot/common/robot_devices/cameras/opencv.py

@@ -80,7 +80,7 @@ def _find_cameras(
     possible_camera_ids: list[int | str], raise_when_empty=False, mock=False
 ) -> list[int | str]:
     if mock:
-        import tests.cameras.mock_cv2 as cv2
+        import lerobot.common.mocks.cameras.mock_cv2 as cv2
     else:
         import cv2
 
@@ -269,7 +269,7 @@ class OpenCVCamera:
         self.logs = {}
 
         if self.mock:
-            import tests.cameras.mock_cv2 as cv2
+            import lerobot.common.mocks.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -286,7 +286,7 @@ class OpenCVCamera:
             raise RobotDeviceAlreadyConnectedError(f"OpenCVCamera({self.camera_index}) is already connected.")
 
         if self.mock:
-            import tests.cameras.mock_cv2 as cv2
+            import lerobot.common.mocks.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -398,7 +398,7 @@ class OpenCVCamera:
         # so we convert the image color from BGR to RGB.
         if requested_color_mode == "rgb":
             if self.mock:
-                import tests.cameras.mock_cv2 as cv2
+                import lerobot.common.mocks.cameras.mock_cv2 as cv2
             else:
                 import cv2
 

lerobot/common/robot_devices/motors/dynamixel.py

@@ -332,7 +332,7 @@ class DynamixelMotorsBus:
             )
 
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -356,7 +356,7 @@ class DynamixelMotorsBus:
 
     def reconnect(self):
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -646,7 +646,7 @@ class DynamixelMotorsBus:
 
     def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -691,7 +691,7 @@ class DynamixelMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -757,7 +757,7 @@ class DynamixelMotorsBus:
 
     def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 
@@ -793,7 +793,7 @@ class DynamixelMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.motors.mock_dynamixel_sdk as dxl
+            import lerobot.common.mocks.motors.mock_dynamixel_sdk as dxl
         else:
             import dynamixel_sdk as dxl
 

lerobot/common/robot_devices/motors/feetech.py

@@ -313,7 +313,7 @@ class FeetechMotorsBus:
             )
 
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -337,7 +337,7 @@ class FeetechMotorsBus:
 
     def reconnect(self):
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -664,7 +664,7 @@ class FeetechMotorsBus:
 
     def read_with_motor_ids(self, motor_models, motor_ids, data_name, num_retry=NUM_READ_RETRY):
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -702,7 +702,7 @@ class FeetechMotorsBus:
 
     def read(self, data_name, motor_names: str | list[str] | None = None):
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -782,7 +782,7 @@ class FeetechMotorsBus:
 
     def write_with_motor_ids(self, motor_models, motor_ids, data_name, values, num_retry=NUM_WRITE_RETRY):
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs
 
@@ -818,7 +818,7 @@ class FeetechMotorsBus:
         start_time = time.perf_counter()
 
         if self.mock:
-            import tests.motors.mock_scservo_sdk as scs
+            import lerobot.common.mocks.motors.mock_scservo_sdk as scs
         else:
             import scservo_sdk as scs