from __future__ import annotations
from collections import OrderedDict
from pathlib import Path
from typing import Any
import numpy as np
import pyrealsense2 as rs
from gymnasium import spaces
from sapien import Pose
from ..utils.camera import pose_CV_ROS, pose_ROS_CV
from ..utils.multiprocessing import SharedObject, ctx, start_and_wait_for_process
from ..utils.realsense import (
RSDevice,
get_connected_rs_devices,
get_default_stream_config,
)
CALIB_CAMERA_POSE_DIR = Path(__file__).resolve().parents[1] / "assets/hec_camera_poses"
CALIB_CAMERA_POSES = {
"front_camera": (
Pose(
np.load(
CALIB_CAMERA_POSE_DIR / "Tb_b2c_20230918_CSE4144_front_jiacheng.npy"
)
)
* pose_CV_ROS
)
}
[docs]
class CameraConfig:
def __init__(
self,
uid: str,
device_sn: str | None = None,
pose: Pose = pose_CV_ROS,
config: tuple[int, int, int] | dict[str, int | tuple[int, int, int]] | None = None, # noqa: E501
*,
preset: str = "Default",
color_option_kwargs={},
depth_option_kwargs={},
json_file: str | Path | None = None,
parent_pose_so_name: str | None = None,
): # fmt: skip
"""Camera configuration.
:param uid: unique id of the camera
:param device_sn: unique serial number of the camera.
If None, use the only RSDevice connected.
:param pose: camera pose in world frame, following ROS frame conventions.
Format is forward(x), left(y) and up(z)
If parent_pose_so_name is not None, this is pose relative to
parent link.
:param config: camera stream config, can be a tuple of (width, height, fps)
or a dict with format {stream_type: (param1, param2, ...)}.
If config is None, use default config realsense.get_default_config().
If config is a tuple, enables color & depth streams with config.
If config is a dict, enables streams given stream parameters.
Possible stream parameters format:
(width, height, fps) for video streams OR
fps for motion_streams
An example config dict:
{"Color": (848, 480, 30), "Depth": (848, 480, 30),
"Infrared 1": (848, 480, 30), "Acceleration": 250}
:param preset: depth sensor presets
:param color_option_kwargs: color sensor optional keywords
:param depth_option_kwargs: depth sensor optional keywords
:param json_file: path to a json file containing sensor configs
:param parent_pose_so_name: SharedObject name of camera's parent link pose
in world frame. Defaults to None (camera has no
mounting parent link).
""" # noqa: E501
self.uid = uid
self.device_sn = device_sn
self.pose = pose
self.config = config
self.preset = preset
self.color_option_kwargs = color_option_kwargs
self.depth_option_kwargs = depth_option_kwargs
self.json_file = json_file
self.parent_pose_so_name = parent_pose_so_name
def __repr__(self) -> str:
return self.__class__.__name__ + "(" + str(self.__dict__) + ")"
[docs]
def parse_camera_cfgs(camera_cfgs):
if isinstance(camera_cfgs, (tuple, list)):
return OrderedDict([(cfg.uid, cfg) for cfg in camera_cfgs])
elif isinstance(camera_cfgs, dict):
return OrderedDict(camera_cfgs)
elif isinstance(camera_cfgs, CameraConfig):
return OrderedDict([(camera_cfgs.uid, camera_cfgs)])
else:
raise TypeError(type(camera_cfgs))
[docs]
def update_camera_cfgs_from_dict(
camera_cfgs: dict[str, CameraConfig], cfg_dict: dict[str, Any | dict[str, Any]]
):
# First, apply global configuration
for k, v in cfg_dict.items():
if k in camera_cfgs: # camera_name, camera-specific config
continue
for cfg in camera_cfgs.values():
if not hasattr(cfg, k):
raise AttributeError(f"{k} is not a valid attribute of CameraConfig")
else:
setattr(cfg, k, v)
# Then, apply camera-specific configuration
for name, v in cfg_dict.items():
if name not in camera_cfgs: # not camera_name, global config
continue
cfg = camera_cfgs[name]
for k in v:
assert hasattr(cfg, k), f"{k} is not a valid attribute of CameraConfig"
cfg.__dict__.update(v)
[docs]
class Camera:
"""Wrapper for RealSense camera (RSDevice)"""
_stream_name2obs_key = {
"Color": "rgb",
"Depth": "depth",
"Infrared 1": "ir_l",
"Infrared 2": "ir_r",
}
def __init__(
self, camera_cfg: CameraConfig, *, record_bag_path: str | Path | None = None
):
"""
:param record_bag_path: path to save bag recording if not None.
Must end with ".bag" if it's a file
"""
self.camera_cfg = camera_cfg
self.uid = camera_cfg.uid
self.device_sn = camera_cfg.device_sn
self.local_pose = camera_cfg.pose
device = get_connected_rs_devices(
self.device_sn
if self.device_sn is not None
else get_connected_rs_devices()[0]
)
self.product_type = device.get_info(rs.camera_info.name).split()[-1]
self.config = camera_cfg.config
if isinstance(self.config, tuple):
self.config = {"Color": self.config, "Depth": self.config}
elif self.config is None:
self.config = get_default_stream_config(self.product_type)
self.config = dict(sorted(self.config.items())) # sort config by stream_name
self.parent_pose_so_name = camera_cfg.parent_pose_so_name
self.record_bag_path = record_bag_path
self.device_proc = ctx.Process(
target=RSDevice,
name=f"RSDevice_{self.uid}",
args=(self.device_sn, self.uid),
kwargs=dict(
config=self.config,
preset=camera_cfg.preset,
color_option_kwargs=camera_cfg.color_option_kwargs,
depth_option_kwargs=camera_cfg.depth_option_kwargs,
json_file=camera_cfg.json_file,
record_bag_path=record_bag_path,
run_as_process=True,
parent_pose_so_name=self.parent_pose_so_name,
local_pose=self.local_pose,
),
)
start_and_wait_for_process(self.device_proc, timeout=30)
# Create SharedObject to control RSDevice and fetch data
self.so_joined = SharedObject(f"join_rs_{self.uid}")
self.so_sync = SharedObject(f"sync_rs_{self.uid}")
self.so_start = SharedObject(f"start_rs_{self.uid}")
self.so_data_dict = {} # {obs_key: SharedObject}
for stream_name in self.config:
self.so_data_dict[self._stream_name2obs_key[stream_name]] = SharedObject(
f"rs_{self.uid}_{stream_name.lower().replace(' ', '_')}"
)
self.so_intr = None
if "rgb" in self.so_data_dict or "depth" in self.so_data_dict: # intrinsics
self.so_intr = SharedObject(f"rs_{self.uid}_intr")
self.so_pose = SharedObject(f"rs_{self.uid}_pose")
self.so_start.assign(True) # start the RSDevice
self.intrinsic_matrix = None
if self.so_intr is not None:
# Wait for intrinsic matrix
while not self.so_intr.modified:
pass
self.intrinsic_matrix = self.so_intr.fetch()
self._camera_buffer = {} # {obs_key: np.ndarray} for self.take_picture()
self._camera_pose = self.so_pose.fetch()
[docs]
def take_picture(self):
"""Fetch images and camera pose from the camera"""
assert self.device_proc.is_alive(), f"RSDevice for {self!r} has died"
# Trigger camera capture in other processes
self.so_sync.trigger()
for obs_key, so_data in self.so_data_dict.items():
if obs_key == "depth":
self._camera_buffer["depth"] = (
so_data.fetch(lambda d: d[..., None].astype(np.float32)) / 1000.0
if self.product_type != "L515"
else 4000.0
)
else:
self._camera_buffer[obs_key] = so_data.fetch()
self._camera_pose = self.so_pose.fetch()
[docs]
def get_images(self, take_picture=False) -> dict[str, np.ndarray]:
"""Get (raw) images from the camera. Takes ~300 us for 848x480 @ 60fps
:return rgb: color image, [H, W, 3] np.uint8 array
:return depth: depth image, [H, W, 1] np.float32 array
:return ir_l: left IR image, [H, W] np.uint8 array
:return ir_r: right IR image, [H, W] np.uint8 array
"""
if take_picture:
self.take_picture()
return self._camera_buffer
@property
def pose(self) -> Pose:
"""Camera pose in world frame, following ROS frame conventions
Format is forward(x), left(y) and up(z)
"""
if self.parent_pose_so_name is not None: # dynamic camera pose
return self.so_pose.fetch()
else: # static camera pose
return self.local_pose
[docs]
def get_extrinsic_matrix(self, pose: Pose = None) -> np.ndarray:
"""Returns a 4x4 extrinsic camera matrix in OpenCV format
right(x), down(y), forward(z)
"""
if pose is not None:
return (pose_CV_ROS * pose.inv()).to_transformation_matrix()
else:
return (pose_CV_ROS * self.pose.inv()).to_transformation_matrix()
[docs]
def get_model_matrix(self, pose: Pose = None) -> np.ndarray:
"""
Returns a 4x4 camera model matrix in OpenCV format
right(x), down(y), forward(z)
Note: this impl is different from sapien where the format is
right(x), up(y), back(z)
"""
if pose is not None:
return (pose * pose_ROS_CV).to_transformation_matrix()
else:
return (self.pose * pose_ROS_CV).to_transformation_matrix()
[docs]
def get_params(self):
"""Get camera parameters."""
pose = self._camera_pose
return dict(
extrinsic_cv=self.get_extrinsic_matrix(pose),
cam2world_cv=self.get_model_matrix(pose),
intrinsic_cv=self.intrinsic_matrix,
)
@property
def observation_space(self) -> spaces.Dict:
obs_spaces = OrderedDict()
for stream_name, params in self.config.items():
if isinstance(params, tuple):
W, H, _ = params
if stream_name == "Color":
shape = (H, W, 3)
elif stream_name == "Depth":
shape = (
self.config["Color"][1::-1] + (1,)
if "Color" in self.config
else (H, W, 1)
)
else:
shape = (H, W)
dtype = np.float32 if stream_name == "Depth" else np.uint8
obs_spaces[self._stream_name2obs_key[stream_name]] = spaces.Box(
low=0,
high=255 if dtype == np.uint8 else np.inf,
shape=shape,
dtype=dtype,
)
else:
raise NotImplementedError(f"No support for {stream_name=} yet")
return spaces.Dict(obs_spaces)
def __del__(self):
self.so_joined.trigger()
self.device_proc.join()
def __repr__(self):
return (
f"<{self.__class__.__name__}: {self.uid} (S/N: {self.device_sn}) "
f"config={self.config}>"
)