Source code for sapien.sensor.simsense_component

from ..pysapien import Component, Pose, CudaArray
from ..pysapien.simsense import DepthSensorEngine

import numpy as np

from typing import Union


[docs] class SimSenseComponent(Component): def __init__( self, rgb_resolution: tuple, ir_resolution: tuple, rgb_intrinsic: np.ndarray, ir_intrinsic: np.ndarray, trans_pose_l: Pose, trans_pose_r: Pose, min_depth: float, max_depth: float, ir_noise_seed: int, ir_speckle_noise: float, ir_thermal_noise: float, rectified: bool, census_width: int, census_height: int, max_disp: int, block_width: int, block_height: int, p1_penalty: int, p2_penalty: int, uniqueness_ratio: int, lr_max_diff: int, median_filter_size: int, depth_dilation: bool, ): super().__init__() # Instance check if not isinstance(rgb_resolution[0], int) or not isinstance( rgb_resolution[1], int ): raise TypeError("RGB resolution (width and height) must be integer") if ( not isinstance(ir_resolution[0], int) or not isinstance(ir_resolution[1], int) or ir_resolution[0] < 32 or ir_resolution[1] < 32 ): raise TypeError( "Infrared resolution (width and height) must be integer and no less than 32" ) if ir_speckle_noise > 0 and ir_thermal_noise <= 0: raise TypeError( "ir_speckle_noise > 0, Infrared noise simulation is on. ir_thermal_noise must also be positive" ) if ( not isinstance(census_width, int) or not isinstance(census_height, int) or census_width <= 0 or census_height <= 0 or census_width % 2 == 0 or census_height % 2 == 0 or census_width * census_height > 65 ): raise TypeError( "census_width and census_height must be positive odd integers and their product should be no larger than 65" ) if not isinstance(max_disp, int) or max_disp < 32 or max_disp > 1024: raise TypeError("max_disp must be integer and within range [32, 1024]") if ( not isinstance(block_width, int) or not isinstance(block_height, int) or block_width <= 0 or block_height <= 0 or block_width % 2 == 0 or block_height % 2 == 0 or block_width * block_height > 256 ): raise TypeError( "block_width and block_height must be positive odd integers and their product should be no larger than 256" ) if ( not isinstance(p1_penalty, int) or not isinstance(p2_penalty, int) or p1_penalty <= 0 or p2_penalty <= 0 or p1_penalty >= p2_penalty or p2_penalty >= 224 ): raise TypeError( "p1_penalty must be positive integer less than p2_penalty and p2_penalty be positive integer less than 224" ) if ( not isinstance(uniqueness_ratio, int) or uniqueness_ratio < 0 or uniqueness_ratio > 255 ): raise TypeError( "uniqueness_ratio must be positive integer and no larger than 255" ) if not isinstance(lr_max_diff, int) or lr_max_diff < -1 or lr_max_diff > 255: raise TypeError("lr_max_diff must be integer and within the range [0, 255]") if ( median_filter_size != 1 and median_filter_size != 3 and median_filter_size != 5 and median_filter_size != 7 ): raise TypeError("Median filter size choices are 1, 3, 5, 7") self.rgb_resolution = rgb_resolution self.ir_resolution = ir_resolution self.rgb_intrinsic = rgb_intrinsic self.ir_intrinsic = ir_intrinsic self.trans_pose_l = trans_pose_l self.trans_pose_r = trans_pose_r self.min_depth = min_depth self.max_depth = max_depth self.ir_noise_seed = ir_noise_seed self.ir_speckle_noise = ir_speckle_noise self.ir_thermal_noise = ir_thermal_noise self.rectified = rectified self.census_width = census_width self.census_height = census_height self.max_disp = max_disp self.block_width = block_width self.block_height = block_height self.p1_penalty = p1_penalty self.p2_penalty = p2_penalty self.uniqueness_ratio = uniqueness_ratio self.lr_max_diff = lr_max_diff self.median_filter_size = median_filter_size self.depth_dilation = depth_dilation self._default_speckle_shape = 1333.33 self._default_gaussian_sigma = 0.25 self._default_gaussian_mu = 0 self._engine = None
[docs] def on_add_to_scene(self, scene): import cv2 if self.ir_speckle_noise == 0: speckle_shape = 0 else: speckle_shape = self._default_speckle_shape / self.ir_speckle_noise speckle_scale = self.ir_speckle_noise / self._default_speckle_shape gaussian_mu = self._default_gaussian_mu gaussian_sigma = self._default_gaussian_sigma * self.ir_thermal_noise # Depth computing engine ir_size, rgb_size = self.ir_resolution, self.rgb_resolution ir_intrinsic = self.ir_intrinsic.astype(float) rgb_intrinsic = self.rgb_intrinsic.astype(float) rgb_pose = Pose() # Dummy pose for deriving extrinsic matrices rgb_extrinsic = self._pose2cv2ex(rgb_pose) l_extrinsic = self._pose2cv2ex(rgb_pose * self.trans_pose_l) r_extrinsic = self._pose2cv2ex(rgb_pose * self.trans_pose_r) l2r = r_extrinsic @ np.linalg.inv(l_extrinsic).astype(float) l2rgb = rgb_extrinsic @ np.linalg.inv(l_extrinsic).astype(float) r1, r2, p1, p2, q, _, _ = cv2.stereoRectify( cameraMatrix1=ir_intrinsic, distCoeffs1=None, cameraMatrix2=ir_intrinsic, distCoeffs2=None, imageSize=ir_size, R=l2r[:3, :3], T=l2r[:3, 3:], alpha=1.0, newImageSize=ir_size, ) f_len = q[2][3] # focal length of the left camera in meters b_len = 1.0 / q[3][2] # baseline length in meters map_l = cv2.initUndistortRectifyMap( ir_intrinsic, None, r1, p1, ir_size, cv2.CV_32F ) map_r = cv2.initUndistortRectifyMap( ir_intrinsic, None, r2, p2, ir_size, cv2.CV_32F ) map_lx, map_ly = map_l map_rx, map_ry = map_r a1, a2, a3, b = self._get_registration_mat( ir_size, ir_intrinsic, rgb_intrinsic, l2rgb ) rgb_fx = rgb_intrinsic[0][0] rgb_fy = rgb_intrinsic[1][1] rgb_skew = rgb_intrinsic[0][1] rgb_cx = rgb_intrinsic[0][2] rgb_cy = rgb_intrinsic[1][2] self._engine = DepthSensorEngine( ir_size[1], ir_size[0], rgb_size[1], rgb_size[0], f_len, b_len, self.min_depth, self.max_depth, self.ir_noise_seed, speckle_shape, speckle_scale, gaussian_mu, gaussian_sigma, self.rectified, self.census_width, self.census_height, self.max_disp, self.block_width, self.block_height, self.p1_penalty, self.p2_penalty, self.uniqueness_ratio, self.lr_max_diff, self.median_filter_size, map_lx, map_ly, map_rx, map_ry, a1, a2, a3, b[0], b[1], b[2], self.depth_dilation, rgb_fx, rgb_fy, rgb_skew, rgb_cx, rgb_cy, )
[docs] def on_remove_from_scene(self, scene): self._engine = None
[docs] def compute( self, left: Union[np.ndarray, CudaArray], right: Union[np.ndarray, CudaArray], bbox_start: tuple = None, bbox_size: tuple = None, ) -> None: if self._engine is None: raise RuntimeError("simsense component is not added to scene") if bbox_start is not None and bbox_size is not None: self._engine.compute(left, right, True, *bbox_start, *bbox_size) else: self._engine.compute(left, right, False, 0, 0, 0, 0)
[docs] def get_ndarray(self) -> np.ndarray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_ndarray()
[docs] def get_cuda(self) -> CudaArray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_cuda()
[docs] def get_point_cloud_ndarray(self) -> np.ndarray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_point_cloud_ndarray()
[docs] def get_point_cloud_cuda(self) -> CudaArray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_point_cloud_cuda()
[docs] def get_rgb_point_cloud_ndarray(self, rgba_cuda: CudaArray) -> np.ndarray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_rgb_point_cloud_ndarray(rgba_cuda)
[docs] def get_rgb_point_cloud_cuda(self, rgba_cuda: CudaArray) -> CudaArray: if self._engine is None: raise RuntimeError("simsense component is not added to scene") return self._engine.get_rgb_point_cloud_cuda(rgba_cuda)
@staticmethod def _get_registration_mat(ir_size, ir_intrinsic, rgb_intrinsic, ir2rgb): R = ir2rgb[:3, :3] t = ir2rgb[:3, 3:] w, h = ir_size x = np.arange(w) y = np.arange(h) u, v = np.meshgrid(x, y) w = np.ones_like(u) pixel_coords = np.stack([u, v, w], axis=-1) # pixel_coords[y, x] is (x, y, 1) A = np.einsum( "ij,hwj->hwi", rgb_intrinsic @ R @ np.linalg.inv(ir_intrinsic), pixel_coords ) B = rgb_intrinsic @ t return A[..., 0], A[..., 1], A[..., 2], B @staticmethod def _pose2cv2ex(pose): ros2opencv = np.array( [ [0.0, -1.0, 0.0, 0.0], [0.0, 0.0, -1.0, 0.0], [1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0], ], dtype=float, ) return ros2opencv @ np.linalg.inv(pose.to_transformation_matrix())