Source code for sapien.sensor.stereodepth

from ..pysapien import Pose, Entity
from ..pysapien.render import (
    RenderCameraComponent,
    RenderTexturedLightComponent,
    RenderTexture2D,
)
from typing import Optional
from copy import deepcopy as copy

import os
import numpy as np


[docs] class StereoDepthSensorConfig: """ An instance of this class is required to initialize StereoDepthSensor. """ def __init__(self): self.light_pattern = os.path.join( os.path.dirname(__file__), "assets/patterns/d415.png" ) """Path to active light pattern file. Use RGB modality if set to None.""" self.ir_camera_exposure = 0.01 # """Camera exposure for infrared cameras.""" self.rgb_resolution = (1920, 1080) """Resolution of the RGB camera (width x height).""" self.ir_resolution = (1280, 720) """Resolution of the infrared cameras (width x height).""" self.rgb_intrinsic = np.array( [[1380.0, 0.0, 960.0], [0.0, 1380.0, 540.0], [0.0, 0.0, 1.0]] ) """Intrinsic matrix of the RGB camera.""" self.ir_intrinsic = np.array( [[920.0, 0.0, 640.0], [0.0, 920.0, 360.0], [0.0, 0.0, 1.0]] ) """Intrinsic matrix of the infrared cameras.""" self.trans_pose_l = Pose([0, -0.0175, 0]) """Relative pose of the left infrared camera to the RGB camera.""" self.trans_pose_r = Pose([0, -0.0720, 0]) """Relative pose of the right infrared camera to the RGB camera.""" self.min_depth = 0.2 """Minimum valid depth in meters.""" self.max_depth = 10.0 """Maximum valid depth (non-inclusive) in meters.""" self.ir_noise_seed = 0 """Random seed for simulating infrared noise.""" self.ir_speckle_noise = 1.0 """Scale for simulating infrared speckle noise. Set to 0 will disable noise simulation.""" self.ir_thermal_noise = 1.0 """Scale for simulating infrared thermal noise. Only effective when speckle noise is on.""" self.rectified = True """Set to true if left and right infrared cameras have a common image plane, otherwise false.""" self.census_width = 7 """Width of the center-symmetric census transform window for steoreo matching. This must be an odd number.""" self.census_height = 7 """Height of the center-symmetric census transform window for stereo matching. This must be an odd number.""" self.max_disp = 128 """Maximum disparity (non-inclusive) for stereo matching.""" self.block_width = 7 """Width of the matched block for stereo matching. This must be an odd number.""" self.block_height = 7 """Height of the matched block for stereo matching. This must be an odd number.""" self.p1_penalty = 8 """P1 penalty for semi-global matching algorithm.""" self.p2_penalty = 32 """P2 penalty for semi-global matching algorithm.""" self.uniqueness_ratio = 15 """How much (in percentage) the minimum cost computed in stereo matching must exceed the second best cost (ignoring best match's adjacent pixels) to be considered valid.""" self.lr_max_diff = 1 """Maximum allowed difference of the left-right consistency check in stereo matching. Set it to 255 will disable the check.""" self.median_filter_size = 3 """Size of the median filter for disparity post-processing. Choices are 1, 3, 5, 7. Set to 1 will disable median filter.""" self.depth_dilation = True """Final depth will be transformed into the same size of rgb_resolution, and depth dilation can dilate the final depth map to avoid holes. Recommended to set as true if rgb_resolution is greater than ir_resolution."""
[docs] class StereoDepthSensor: """ This class simulates an active stereo depth sensor. It has one RGB camera and two infrared camera. Depth is computed via semi-global block matching. Refer to StereoDepthSensorConfig for configurable parameters. The computed depth map will be presented in RGB camera frame. """ def __init__( self, config: StereoDepthSensorConfig, mount_entity: Entity, pose: Optional[Pose] = None, ): """ :param config: configuration of the sensor. :param mount_entity: entity that the sensor is mounted to. :param pose: local pose relative to the mounted entity. If not given, the relative trasformation will be identity transformation. """ super().__init__() from .simsense_component import SimSenseComponent # Basic configuration self._config = config self._mount = mount_entity if pose is None: self._pose = Pose() else: self._pose = pose # Cameras self._cam_rgb = None self._cam_ir_l = None self._cam_ir_r = None self._create_cameras() # Active light self._alight = None self._create_light() # Simsense component self._ss = SimSenseComponent( config.rgb_resolution, config.ir_resolution, config.rgb_intrinsic, config.ir_intrinsic, config.trans_pose_l, config.trans_pose_r, config.min_depth, config.max_depth, config.ir_noise_seed, config.ir_speckle_noise, config.ir_thermal_noise, config.rectified, config.census_width, config.census_height, config.max_disp, config.block_width, config.block_height, config.p1_penalty, config.p2_penalty, config.uniqueness_ratio, config.lr_max_diff, config.median_filter_size, config.depth_dilation, ) self._mount.add_component(self._ss)
[docs] def take_picture(self, infrared_only: bool = False): """ Note: We expect one scene.update_render() call before calling take_picture(). :param infrared_only: If true, only take infrared pictures without taking RGB picture. """ scene = self._mount.get_scene() if scene is None: raise RuntimeError("Cannot take picture: sensor is not added to scene") if not infrared_only: self._cam_rgb.take_picture() self._ir_mode() scene.update_render() self._cam_ir_l.take_picture() self._cam_ir_r.take_picture() self._normal_mode() scene.update_render()
[docs] def compute_depth(self, bbox_start: tuple = None, bbox_size: tuple = None): left_cuda = self._cam_ir_l.get_picture_cuda("Color") right_cuda = self._cam_ir_r.get_picture_cuda("Color") self._ss.compute(left_cuda, right_cuda, bbox_start, bbox_size)
[docs] def set_local_pose(self, pose: Pose): """ Set local pose of the sensor relative to mounted actor. """ self._pose = pose self._alight.local_pose = pose self._cam_rgb.local_pose = pose self._cam_ir_l.local_pose = pose * self._config.trans_pose_l self._cam_ir_r.local_pose = pose * self._config.trans_pose_r
[docs] def set_ir_noise(self, ir_speckle_noise: float, ir_thermal_noise: float): """ :param ir_speckle_noise: Scale for simulating infrared speckle noise. Set to 0 will disable noise simulation. :param ir_thermal_noise: Scale for simulating infrared thermal noise. Only effective when speckle noise is on. """ 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 ir_speckle_noise == 0: speckle_shape = 0 else: speckle_shape = self._default_speckle_shape / ir_speckle_noise speckle_scale = ir_speckle_noise / self._default_speckle_shape gaussian_mu = self._default_gaussian_mu gaussian_sigma = self._default_gaussian_sigma * ir_thermal_noise self._config.ir_speckle_noise = ir_speckle_noise self._config.ir_thermal_noise = ir_thermal_noise self._ss._engine.set_ir_noise_parameters( speckle_shape, speckle_scale, gaussian_mu, gaussian_sigma )
[docs] def set_census_window_size(self, census_width: int, census_height: int): """ :param census_width: Width of the center-symmetric census transform window. This must be an odd number. :param census_height: Height of the center-symmetric census transform window. This must be an odd number. """ 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" ) self._config.census_width = census_width self._config.census_height = census_height self._ss._engine.set_census_window_size(census_width, census_height)
[docs] def set_matching_block_size(self, block_width: int, block_height: int): """ :param block_width: Width of the matched block. This must be an odd number. :param block_height: Height of the matched block. This must be an odd number. """ 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" ) self._config.block_width = block_width self._config.block_height = block_height self._ss._engine.set_matching_block_size(block_width, block_height)
[docs] def set_penalties(self, p1_penalty: int, p2_penalty: int): """ :param p1_penalty: P1 penalty for semi-global matching algorithm. :param p2_penalty: P2 penalty for semi-global matching algorithm. """ 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 must be positive integer less than p2 and p2 be positive integer less than 224" ) self._config.p1_penalty = p1_penalty self._config.p2_penalty = p2_penalty self._ss._engine.set_penalties(p1_penalty, p2_penalty)
[docs] def set_uniqueness_ratio(self, uniqueness_ratio: int): """ :param uniqueness_ratio: Margin in percentage by which the minimum computed cost should win the second best (not considering best match's adjacent pixels) cost to consider the found match valid. """ if ( not isinstance(uniqueness_ratio, int) or uniqueness_ratio < 0 or uniqueness_ratio > 255 ): raise TypeError( "uniqueness_ratio must be positive integer no larger than 255" ) self._config.uniqueness_ratio = uniqueness_ratio self._ss._engine.set_uniqueness_ratio(uniqueness_ratio)
[docs] def set_lr_max_diff(self, lr_max_diff: int): """ :param lr_max_diff: Maximum allowed difference in the left-right consistency check. Set it to 255 to disable the check. """ 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 within the range [0, 255]") self._config.lr_max_diff = lr_max_diff self._ss._engine.set_lr_max_diff(lr_max_diff)
[docs] def get_config(self): return copy(self._config)
[docs] def get_pose(self): return copy(self._mount.get_pose() * self._pose)
[docs] def get_rgb(self): rgb = self._cam_rgb.get_picture("Color")[..., :3].clip(0, 1) return copy(rgb)
[docs] def get_rgba_cuda(self): return self._cam_rgb.get_picture_cuda("Color")
[docs] def get_ir(self): """ Note: Noise simulation won't be reflected here. """ ir_l = self._cam_ir_l.get_picture("Color")[..., 0].clip(0, 1) ir_r = self._cam_ir_r.get_picture("Color")[..., 0].clip(0, 1) return [copy(ir_l), copy(ir_r)]
[docs] def get_depth(self): """ Note: Returned depth map will be of the same resolution and frame of RGB camera. """ depth = self._ss.get_ndarray() return copy(depth)
[docs] def get_depth_cuda(self): """ Note: Returned depth map will be of the same resolution and frame of RGB camera. """ return self._ss.get_cuda()
[docs] def get_pointcloud(self, with_rgb: bool = False): """ Note: Returned point cloud is from RGB camera's with x rightward, y downward, z forward. """ if with_rgb: rgba_cuda = self.get_rgba_cuda() pc = self._ss.get_rgb_point_cloud_ndarray(rgba_cuda) else: pc = self._ss.get_point_cloud_ndarray() return pc
[docs] def get_pointcloud_cuda(self, with_rgb: bool = False): """ Note: Returned point cloud is from RGB camera's with x rightward, y downward, z forward. """ if with_rgb: rgba_cuda = self.get_rgba_cuda() pc_cuda = self._ss.get_rgb_point_cloud_cuda(rgba_cuda) else: pc_cuda = self._ss.get_point_cloud_cuda() return pc_cuda
def _create_cameras(self): self._cam_rgb = RenderCameraComponent(*self._config.rgb_resolution) self._cam_rgb.local_pose = self._pose self._cam_rgb.name = "cam_rgb" self._cam_rgb.set_perspective_parameters( 0.1, 100.0, self._config.rgb_intrinsic[0, 0], self._config.rgb_intrinsic[1, 1], self._config.rgb_intrinsic[0, 2], self._config.rgb_intrinsic[1, 2], self._config.rgb_intrinsic[0, 1], ) self._cam_ir_l = RenderCameraComponent(*self._config.ir_resolution) self._cam_ir_l.local_pose = self._pose * self._config.trans_pose_l self._cam_ir_l.name = "cam_ir_l" self._cam_ir_l.set_perspective_parameters( 0.1, 100.0, self._config.ir_intrinsic[0, 0], self._config.ir_intrinsic[1, 1], self._config.ir_intrinsic[0, 2], self._config.ir_intrinsic[1, 2], self._config.ir_intrinsic[0, 1], ) self._cam_ir_r = RenderCameraComponent(*self._config.ir_resolution) self._cam_ir_r.local_pose = self._pose * self._config.trans_pose_r self._cam_ir_r.name = "cam_ir_r" self._cam_ir_r.set_perspective_parameters( 0.1, 100.0, self._config.ir_intrinsic[0, 0], self._config.ir_intrinsic[1, 1], self._config.ir_intrinsic[0, 2], self._config.ir_intrinsic[1, 2], self._config.ir_intrinsic[0, 1], ) self._mount.add_component(self._cam_rgb) self._mount.add_component(self._cam_ir_l) self._mount.add_component(self._cam_ir_r) self._cam_ir_l.set_property("exposure", float(self._config.ir_camera_exposure)) self._cam_ir_r.set_property("exposure", float(self._config.ir_camera_exposure)) def _create_light(self): # Active Light self._alight = RenderTexturedLightComponent() self._alight.color = [0, 0, 0] self._alight.inner_fov = 1.57 self._alight.outer_fov = 1.57 self._alight.texture = RenderTexture2D(self._config.light_pattern) self._alight.local_pose = self._pose self._alight.name = "active_light" self._mount.add_component(self._alight) def _ir_mode(self): if self._config.light_pattern is None: return self._alight.color = [100, 0, 0] def _normal_mode(self): if self._config.light_pattern is None: return self._alight.color = [0, 0, 0]