Support set_local_pose in StereoDepthSensor (#112)

* Support local pose for StereoDepthSensor

* support set_local_pose in StereoDepthSensor

* minor fix

---------

Co-authored-by: Jiayuan-Gu <jigu@eng.ucsd.edu>

Author: angli66

docs/source/tutorial/rendering/depth_sensor.rst

@@ -67,18 +67,18 @@ Optionally, you can mount the sensor to an existing actor:
 ::
 
     sensor_config = StereoDepthSensorConfig()
-    sensor = StereoDepthSensor('sensor', scene, sensor_config, mount=actor)
+    sensor = StereoDepthSensor('sensor', scene, sensor_config, mount=actor, pose=Pose()) # pose is relative to mount
 
-After mounting to an actor, the sensor will move along with it. Calling ``set_pose`` from either the sensor or the actor will move the two entities together.
+After mounting to an actor, the sensor will move along with it. Calling ``sensor.set_local_pose`` can adjust the relative pose to the mounted actor.
 
-``sensor`` behaves very similar to a camera. You can ``set_pose`` and ``take_picture`` just like working with a camera. What's more, you can ``compute_depth`` and
-``get_pointcloud`` on the sensor:
+``sensor`` behaves very similar to a camera. You can ``set_pose`` (or ``set_local_pose`` when mounted) and ``take_picture`` just like working with a camera.
+In addtion to these basic functions, you can ``compute_depth`` and ``get_pointcloud`` on the sensor:
 
 .. literalinclude:: ../../../../examples/rendering/rt_stereodepth.py
     :dedent: 0
     :lines: 90-95
 
-One important differences between camera and ``sensor`` is that while camera will only take picture of an RGB image, ``sensor`` will take another pair
+One important difference between camera and ``sensor`` is that while camera will only take picture of an RGB image, ``sensor`` will take another pair
 of infrared images, which will be used to compute depth. After calling ``take_picture``, the RGB image and infrared images will be saved within ``sensor``.
 Calling ``sensor.get_rgb`` and ``sensor.get_ir`` will return the pictures in ndarray form. Let's take a look at them:
 
@@ -144,13 +144,14 @@ If depth is the only needed data and RGB data is not needed, you can specify ``i
 This can save the time for rendering RGB image.
 
 
-As mentioned above, the final depth map generated by ``StereoDepthSensor`` will be transformed into RGB camera frame. It will be of the same resolution and frame as of the RGB
-camera. This feature can be used to achieve fast downsampling/upsampling. All you need is to specify the ``rgb_resolution`` of the ``StereoDepthSensorConfig`` as the final
+As mentioned above, the final depth map generated by ``StereoDepthSensor`` will be transformed into RGB camera frame. It will be of the same resolution and frame as of the
+RGB camera. This feature can be used to achieve fast downsampling/upsampling. All you need is to specify the ``rgb_resolution`` of the ``StereoDepthSensorConfig`` as the final
 resolution you want to sample on, and ``StereoDepthSensor`` will do that for you on the GPU. In this way, you don't need to attach any slow CPU resampling function to it.
 
-In general, lowering ``rt_samples_per_pixel`` for renderer, ``ir_resolution`` of sensor config and ``max_disp`` of sensor config are all good ways to enhance  computation
-speed. However, note that these changes might also have effect on the output depth map. You can freely adjust the parameters until you find the settings that satisfies
-your need.
+In general, lowering ``rt_samples_per_pixel`` for renderer, ``ir_resolution`` of sensor config and ``max_disp`` of sensor config are all good ways to enhance computation
+speed. For example, try changing the ``rt_samples_per_pixel`` in our script from ``32`` to ``8``. You might find that we can get almost as good results as before while
+``sensor.take_picture`` now runs much faster. However, note that these changes generally will have some effect (trade quality with speed) on the output depth map. You can freely
+adjust the parameters until you find the settings that satisfies your need.
 
 
 Seamless transfer to other GPU libraries

examples/rendering/rt_stereodepth.py

@@ -87,7 +87,7 @@ def main():
 
     sensor_config = StereoDepthSensorConfig()
     sensor = StereoDepthSensor('sensor', scene, sensor_config)
-    sensor.set_pose(Pose([-0.13732728, 0.13584249, 0.82386769],[0.84667092, -0.01287458, 0.53195302, -0.00292851]))
+    sensor.set_pose(Pose([-0.13732728, 0.13584249, 0.82386769], [0.84667092, -0.01287458, 0.53195302, -0.00292851]))
 
     scene.step()
     scene.update_render()

manualtest/rt_stereodepth.py

@@ -87,7 +87,7 @@ def main():
 
     sensor_config = StereoDepthSensorConfig()
     sensor = StereoDepthSensor('sensor', scene, sensor_config)
-    sensor.set_pose(Pose([-0.13732728, 0.13584249, 0.82386769],[0.84667092, -0.01287458, 0.53195302, -0.00292851]))
+    sensor.set_pose(Pose([-0.13732728, 0.13584249, 0.82386769], [0.84667092, -0.01287458, 0.53195302, -0.00292851]))
 
     scene.step()
     scene.update_render()

manualtest/rt_stereodepth2.py

@@ -102,8 +102,7 @@ def main():
     ])
     sensor_config.max_disp = 96 # Max disparity for stereo matching
     sensor = StereoDepthSensor('sensor', scene, sensor_config)
-    sensor.set_pose(
-        Pose([-0.28, -0.28, 0.46], [0.8876263, -0.135299, 0.3266407, 0.2951603]))
+    sensor.set_pose(Pose([-0.28, -0.28, 0.46], [0.8876263, -0.135299, 0.3266407, 0.2951603]))
 
     # Warm up
     scene.step()

python/py_package/sensor/stereodepth.py

@@ -121,7 +121,15 @@ class StereoDepthSensor(SensorEntity):
     parameters. The computed depth map will be presented in RGB camera frame.
     """
 
-    def __init__(self, sensor_name: str, scene: Scene, config: StereoDepthSensorConfig, mount: Optional[Actor] = None):
+    def __init__(self, sensor_name: str, scene: Scene, config: StereoDepthSensorConfig, mount: Optional[Actor] = None, pose: Optional[Pose] = None):
+        """
+        :param sensor_name: name of the sensor.
+        :param scene: scene that the sensor is attached to.
+        :param config: configuration of the sensor.
+        :param mount: optionally mount the sensor to an actor. If mounted, the sensor will move along with the actor.
+        :param pose: If not mounted, this will be the global pose for the sensor. Otherwise, it will be the local pose relative to the mounted actor.
+        """
+  
         super().__init__()
 
         # Basic configuration
@@ -169,15 +177,19 @@ def __init__(self, sensor_name: str, scene: Scene, config: StereoDepthSensorConf
             self._config.median_filter_size != 7:
             raise TypeError("Median filter size choices are 1, 3, 5, 7")
 
-        # Pose
+        # self._pose is global if not mounted and local if mounted
         self._mount = mount
-        self._alight = self._scene.add_active_light(
-            pose=Pose(), color=[0, 0, 0], fov=1.57, tex_path=self._config.light_pattern)
-        if self._mount is None:
+        if pose is None:
             self._pose = Pose()
         else:
+            self._pose = pose
+
+        # Active Light
+        self._alight = self._scene.add_active_light(
+            pose=Pose(), color=[0, 0, 0], fov=1.57, tex_path=self._config.light_pattern)
+        if self._mount is not None:
             self._alight.set_parent(self._mount, keep_pose=False)
-            self._alight.set_local_pose(Pose())
+        self._alight.set_local_pose(self._pose)
 
         # Cameras
         self._cam_rgb = None
@@ -201,10 +213,7 @@ def __init__(self, sensor_name: str, scene: Scene, config: StereoDepthSensorConf
         ir_size, rgb_size = self._config.ir_resolution, self._config.rgb_resolution
         ir_intrinsic = self._config.ir_intrinsic.astype(float)
         rgb_intrinsic = self._config.rgb_intrinsic.astype(float)
-        if self._mount is None:
-            rgb_pose = self._pose
-        else:
-            rgb_pose = self._mount.get_pose()
+        rgb_pose = self._cam_rgb.get_pose()
         rgb_extrinsic = self._pose2cv2ex(rgb_pose)
         l_extrinsic = self._pose2cv2ex(rgb_pose * self._config.trans_pose_l)
         r_extrinsic = self._pose2cv2ex(rgb_pose * self._config.trans_pose_r)
@@ -258,16 +267,26 @@ def compute_depth(self):
 
     def set_pose(self, pose: Pose):
         """
-        Note: If mount exists, setting sensor's pose will also have effect on the mounted actor.
+        Set the global pose of the sensor. Call set_local_pose instead of set_pose if mount exists.
         """
         if self._mount is None:
             self._pose = pose
-            self._alight.set_pose(self._pose)
-            self._cam_rgb.set_local_pose(self._pose)
-            self._cam_ir_l.set_local_pose(self._pose * self._config.trans_pose_l)
-            self._cam_ir_r.set_local_pose(self._pose * self._config.trans_pose_r)
+            self._alight.set_pose(pose)
+            self._cam_rgb.set_local_pose(pose)
+            self._cam_ir_l.set_local_pose(pose * self._config.trans_pose_l)
+            self._cam_ir_r.set_local_pose(pose * self._config.trans_pose_r)
         else:
-            self._mount.set_pose(pose)
+            raise RuntimeError("cannot set_pose when sensor is mounted, use set_local_pose instead")
+
+    def set_local_pose(self, pose: Pose):
+        """
+        If mount exists, set local pose of the sensor relative to mounted actor. Otherwise, set global pose of the sensor.
+        """
+        self._pose = pose
+        self._alight.set_local_pose(pose)
+        self._cam_rgb.set_local_pose(pose)
+        self._cam_ir_l.set_local_pose(pose * self._config.trans_pose_l)
+        self._cam_ir_r.set_local_pose(pose * self._config.trans_pose_r)
 
     def set_ir_noise(self, ir_speckle_noise: float, ir_thermal_noise: float):
         """
@@ -350,7 +369,7 @@ def get_pose(self):
         if self._mount is None:
             return copy(self._pose)
         else:
-            return copy(self._mount().get_pose())
+            return copy(self._mount.get_pose() * self._pose)
 
     def get_rgb(self):
         rgb = self._cam_rgb.get_color_rgba()[..., :3].clip(0, 1)
@@ -421,13 +440,13 @@ def _create_cameras(self):
             self._cam_ir_r.set_local_pose(self._pose * self._config.trans_pose_r)
         else:
             self._cam_rgb = self._scene.add_mounted_camera(
-                f"{self.name}", self._mount, Pose(), *self._config.rgb_resolution, 0.78, 0.001, 100)
+                f"{self.name}", self._mount, self._pose, *self._config.rgb_resolution, 0.78, 0.001, 100)
 
             self._cam_ir_l = self._scene.add_mounted_camera(
-                f"{self.name}_left", self._mount, self._config.trans_pose_l, *self._config.ir_resolution, 0.78, 0.001, 100)
+                f"{self.name}_left", self._mount, self._pose * self._config.trans_pose_l, *self._config.ir_resolution, 0.78, 0.001, 100)
 
             self._cam_ir_r = self._scene.add_mounted_camera(
-                f"{self.name}_right", self._mount, self._config.trans_pose_r, *self._config.ir_resolution, 0.78, 0.001, 100)
+                f"{self.name}_right", self._mount, self._pose * self._config.trans_pose_r, *self._config.ir_resolution, 0.78, 0.001, 100)
 
         self._cam_rgb.set_perspective_parameters(
             0.1, 100.0,