style: fix python linting

Author: juelg

extensions/rcs_realsense/src/rcs_realsense/calibration.py

@@ -21,7 +21,9 @@ def __init__(self, camera_name: str):
         # base frame to camera, world to base frame
         self._extrinsics: np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None = None
         self.camera_name = camera_name
-        self.tag_to_world = common.Pose(rpy_vector=[np.pi, 0, -np.pi / 2], translation=[0.145, 0, 0]).pose_matrix()
+        self.tag_to_world = common.Pose(
+            rpy_vector=np.array([np.pi, 0, -np.pi / 2]), translation=np.array([0.145, 0, 0])
+        ).pose_matrix()
 
     def calibrate(
         self,
@@ -46,13 +48,11 @@ def calibrate(
                 frames.append(sample.camera.color.data.copy())
         print(frames)
 
-        last_frame, tag_to_cam = get_average_marker_pose(
-            frames, intrinsics=intrinsics, calib_tag_id=9, show_live_window=True
-        )
+        _, tag_to_cam = get_average_marker_pose(frames, intrinsics=intrinsics, calib_tag_id=9, show_live_window=False)
 
         cam_to_world = self.tag_to_world @ np.linalg.inv(tag_to_cam)
         world_to_cam = np.linalg.inv(cam_to_world)
-        self._extrinsics = world_to_cam
+        self._extrinsics = world_to_cam  # type: ignore
         return True
 
     def get_extrinsics(self) -> np.ndarray[tuple[typing.Literal[4], typing.Literal[4]], np.dtype[np.float64]] | None:
@@ -82,8 +82,7 @@ def get_average_marker_pose(
             continue
 
         for corner in marker_det.corners:
-            corner = corner.astype(int)
-            cv2.circle(frame, tuple(corner), 5, (0, 0, 255), -1)
+            cv2.circle(frame, tuple(corner.astype(int)), 5, (0, 0, 255), -1)
 
         poses.append(pose)
 
@@ -92,7 +91,7 @@ def get_average_marker_pose(
         camera_matrix = intrinsics[:3, :3]
 
         if show_live_window:
-            cv2.drawFrameAxes(frame, camera_matrix, None, pose[:3, :3], pose[:3, 3], 0.1)
+            cv2.drawFrameAxes(frame, camera_matrix, None, pose[:3, :3], pose[:3, 3], 0.1)  # type: ignore
             # show frame
             cv2.imshow("frame", frame)
 
@@ -107,12 +106,9 @@ def get_average_marker_pose(
         cv2.destroyAllWindows()
 
     # calculate the average marker pose
-    poses = np.array(poses)
     avg_pose = np.mean(poses, axis=0)
     logger.info(f"Average pose: {avg_pose}")
 
-    # paint avg pose on last frame
-    # cv2.drawFrameAxes(last_frame, camera_matrix, None, avg_pose[:3, :3], avg_pose[:3, 3], 0.1)  # type: ignore
     return last_frame, avg_pose
 
 

extensions/rcs_realsense/src/rcs_realsense/camera.py

@@ -48,7 +48,7 @@ def __init__(
         self.enable_imu = enable_imu
         self.cameras = cameras
         if calibration_strategy is None:
-            calibration_strategy = {camera_name: DummyCalibrationStrategy() for camera_name in cameras.keys()}
+            calibration_strategy = {camera_name: DummyCalibrationStrategy() for camera_name in cameras}
         self.calibration_strategy = calibration_strategy
         self._logger = logging.getLogger(__name__)
         assert (
@@ -123,16 +123,6 @@ def open(self):
     def update_available_devices(self):
         self._available_devices = self.enumerate_connected_devices(rs.context())
 
-    def enable_all_devices(self, enable_ir_emitter: bool = False):
-        """
-        Enable all the Intel RealSense Devices which are connected to the PC
-
-        """
-        self._logger.debug("%i devices have been found", len(self._available_devices))
-
-        for device_info in self._available_devices.values():
-            self.enable_device(device_info, enable_ir_emitter)
-
     def enable_devices(self, devices_to_enable: dict[str, str], enable_ir_emitter: bool = False):
         """
         Enable the Intel RealSense Devices which are connected to the PC
@@ -209,7 +199,7 @@ def enable_device(self, camera_name: str, device_info: RealSenseDeviceInfo, enab
             color_intrinsics=color_intrinsics,
             depth_intrinsics=depth_intrinsics,
             depth_to_color=common.Pose(
-                translation=depth_to_color.translation, rotation=np.array(depth_to_color.rotation).reshape(3, 3)
+                translation=depth_to_color.translation, rotation=np.array(depth_to_color.rotation).reshape(3, 3)  # type: ignore
             ),
         )
 
@@ -266,6 +256,7 @@ def to_ts(frame: rs.frame) -> float:
 
         color_extrinsics = self.calibration_strategy[camera_name].get_extrinsics()
         depth_to_color = device.depth_to_color
+        assert depth_to_color is not None, "Depth to color extrinsics not found"
         depth_extrinsics = (
             color_extrinsics @ depth_to_color.inverse().pose_matrix() if color_extrinsics is not None else None
         )
@@ -285,6 +276,7 @@ def to_ts(frame: rs.frame) -> float:
                 )
             elif rs.stream.depth == stream.stream_type():
                 frame = frameset.get_depth_frame()
+                assert device.depth_scale is not None, "Depth scale not found"
                 depth = DataFrame(
                     data=(to_numpy(frame).astype(np.float64) * device.depth_scale * BaseCameraSet.DEPTH_SCALE).astype(
                         np.uint16
@@ -345,8 +337,10 @@ def enable_emitter(self, enable_ir_emitter=True):
 
     def calibrate(self) -> bool:
         for camera_name in self.cameras:
+            color_intrinsics = self._enabled_devices[camera_name].color_intrinsics
+            assert color_intrinsics is not None, f"Color intrinsics for camera {camera_name} not found"
             if not self.calibration_strategy[camera_name].calibrate(
-                intrinsics=self._enabled_devices[camera_name].color_intrinsics,
+                intrinsics=color_intrinsics,
                 samples=self._frame_buffer[camera_name],
                 lock=self._frame_buffer_lock[camera_name],
             ):

extensions/rcs_realsense/src/rcs_realsense/utils.py

@@ -1,3 +1,6 @@
+import typing
+
+from rcs.camera.hw import CalibrationStrategy
 from rcs_realsense.calibration import FR3BaseArucoCalibration
 from rcs_realsense.camera import RealSenseCameraSet
 
@@ -11,5 +14,5 @@ def default_realsense(name2id: dict[str, str] | None) -> RealSenseCameraSet | No
         name: common.BaseCameraConfig(identifier=id, resolution_width=1280, resolution_height=720, frame_rate=30)
         for name, id in name2id.items()
     }
-    calibration_strategy = {name: FR3BaseArucoCalibration(name) for name in name2id}
+    calibration_strategy = {name: typing.cast(CalibrationStrategy, FR3BaseArucoCalibration(name)) for name in name2id}
     return RealSenseCameraSet(cameras=cameras, calibration_strategy=calibration_strategy)

extensions/rcs_so101/src/rcs_so101/creators.py

@@ -1,4 +1,5 @@
 import logging
+import typing
 from os import PathLike
 from pathlib import Path
 from typing import Type
@@ -15,6 +16,7 @@
 from lerobot.common.teleoperators.so101_leader.so101_leader import SO101Leader
 from lerobot.common.teleoperators.utils import make_teleoperator_from_config
 from rcs.camera.hw import HardwareCameraSet
+from rcs.camera.interface import BaseCameraSet
 from rcs.camera.sim import SimCameraSet
 from rcs.envs.base import (
     CameraSetWrapper,
@@ -125,7 +127,9 @@ def __call__(  # type: ignore
         env = RobotSimWrapper(env, simulation, sim_wrapper)
 
         if cameras is not None:
-            camera_set = SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
+            camera_set = typing.cast(
+                BaseCameraSet, SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
+            )
             env = CameraSetWrapper(env, camera_set, include_depth=True)
 
         if gripper_cfg is not None and isinstance(gripper_cfg, SimGripperConfig):

pyproject.toml

@@ -101,6 +101,7 @@ ignore = [
   "PT018",   # assertion should be broken down into multiple parts
   "NPY002",
   "G004",   # Logging format string does not contain any placeholders
+  "ARG002",
 ]
 
 [tool.pylint.format]

python/rcs/camera/digit_cam.py

@@ -51,3 +51,7 @@ def close(self):
 
     def config(self, camera_name) -> BaseCameraConfig:
         return self.cameras[camera_name]
+
+    def calibrate(self) -> bool:
+        """No calibration needed for DIGIT cameras."""
+        return True

python/rcs/camera/interface.py

@@ -50,7 +50,7 @@ class FrameSet:
 class BaseCameraSet(Protocol):
     """Interface for a set of cameras for sim and hardware"""
 
-    DEPTH_SCALE = 1000
+    DEPTH_SCALE: int = 1000
 
     def buffer_size(self) -> int:
         """Returns size of the internal buffer."""

python/rcs/camera/sim.py

@@ -90,7 +90,7 @@ def _cpp_to_python_frames(self, cpp_frameset: _FrameSet | None) -> FrameSet | No
             frames[color_name] = frame
         return FrameSet(frames=frames, avg_timestamp=cpp_frameset.timestamp)
 
-    def _intrinsics(self, camera_name) -> np.ndarray[tuple[Literal[3, 4]], np.dtype[np.float64]]:
+    def _intrinsics(self, camera_name) -> np.ndarray[tuple[Literal[3], Literal[4]], np.dtype[np.float64]]:
         cam_id = mujoco.mj_name2id(self._sim.model, mujoco.mjtObj.mjOBJ_CAMERA, self.cameras[camera_name].identifier)
         fovy = self._sim.model.cam_fovy[cam_id]
         fx = fy = 0.5 * self.cameras[camera_name].resolution_height / np.tan(fovy * np.pi / 360)
@@ -102,14 +102,14 @@ def _intrinsics(self, camera_name) -> np.ndarray[tuple[Literal[3, 4]], np.dtype[
             ]
         )
 
-    def _extrinsics(self, camera_name) -> np.ndarray[tuple[Literal[4, 4]], np.dtype[np.float64]]:
+    def _extrinsics(self, camera_name) -> np.ndarray[tuple[Literal[4], Literal[4]], np.dtype[np.float64]]:
         cam_id = mujoco.mj_name2id(self._sim.model, mujoco.mjtObj.mjOBJ_CAMERA, self.cameras[camera_name].identifier)
         xpos = self._sim.data.cam_xpos[cam_id]
         xmat = self._sim.data.cam_xmat[cam_id].reshape(3, 3)
 
         cam = common.Pose(rotation=xmat, translation=xpos)
         # put z axis infront
-        rotation_p = common.Pose(rpy_vector=[np.pi, 0, 0], translation=[0, 0, 0])
+        rotation_p = common.Pose(rpy_vector=np.array([np.pi, 0, 0]), translation=np.array([0, 0, 0]))
         cam = cam * rotation_p
 
         return cam.inverse().pose_matrix()

python/rcs/envs/base.py

@@ -3,7 +3,7 @@
 import copy
 import logging
 from enum import Enum, auto
-from typing import Annotated, Any, TypeAlias, cast
+from typing import Annotated, Any, Literal, TypeAlias, cast
 
 import gymnasium as gym
 import numpy as np
@@ -131,7 +131,7 @@ class CameraDataDictType(RCSpaceType):
         "frame",
     ]
     intrinsics: Annotated[
-        np.ndarray,
+        np.ndarray[tuple[Literal[3], Literal[4]], np.dtype[np.float64]] | None,
         gym.spaces.Box(
             low=-np.inf,
             high=np.inf,
@@ -140,7 +140,7 @@ class CameraDataDictType(RCSpaceType):
         ),
     ]
     extrinsics: Annotated[
-        np.ndarray,
+        np.ndarray[tuple[Literal[4], Literal[4]], np.dtype[np.float64]] | None,
         gym.spaces.Box(
             low=-np.inf,
             high=np.inf,
@@ -646,7 +646,7 @@ def check_depth(depth):
                 raise ValueError(msg)
             return self.include_depth
 
-        frame_dict: dict[str, dict[str, np.ndarray]] = {
+        frame_dict: dict[str, dict[str, CameraDataDictType]] = {
             camera_name: (
                 {
                     self.RGB_KEY: CameraDataDictType(

python/rcs/envs/creators.py

@@ -1,11 +1,13 @@
 import logging
+import typing
 from os import PathLike
 from typing import Type
 
 import gymnasium as gym
 import numpy as np
 from gymnasium.envs.registration import EnvCreator
 from rcs._core.sim import CameraType
+from rcs.camera.interface import BaseCameraSet
 from rcs.camera.sim import SimCameraConfig, SimCameraSet
 from rcs.envs.base import (
     CameraSetWrapper,
@@ -87,7 +89,9 @@ def __call__(  # type: ignore
         env = RobotSimWrapper(env, simulation, sim_wrapper)
 
         if cameras is not None:
-            camera_set = SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
+            camera_set = typing.cast(
+                BaseCameraSet, SimCameraSet(simulation, cameras, physical_units=True, render_on_demand=True)
+            )
             env = CameraSetWrapper(env, camera_set, include_depth=True)
 
         if gripper_cfg is not None and isinstance(gripper_cfg, rcs.sim.SimGripperConfig):