[python/viewer] Add support of user-specified additional cameras (rgb…

… and depth map).

python/jiminy_py/src/jiminy_py/viewer/panda3d/panda3d_visualizer.py

@@ -35,7 +35,7 @@
 from direct.gui.OnscreenText import OnscreenText
 from panda3d.core import (  # pylint: disable=no-name-in-module
     NodePath, Point3, Vec3, Vec4, Mat4, Quat, LQuaternion, Geom, GeomEnums,
-    GeomNode, GeomTriangles, GeomVertexData, GeomVertexArrayFormat,
+    GeomNode, GeomTriangles, GeomVertexData, GeomVertexArrayFormat, BitMask32,
     GeomVertexFormat, GeomVertexWriter, PNMImage, PNMImageHeader, TextNode,
     OmniBoundingVolume, CompassEffect, BillboardEffect, InternalName, Filename,
     Material, Texture, TextureStage, TransparencyAttrib, PGTop, Camera, Lens,
@@ -58,7 +58,7 @@
 
 
 WINDOW_SIZE_DEFAULT = (600, 600)
-CAMERA_POS_DEFAULT = [(4.0, -4.0, 1.5), (0, 0, 0.5)]
+CAMERA_POSE_DEFAULT = [(4.0, -4.0, 1.5), (0, 0, 0.5)]
 
 SKY_TOP_COLOR = (0.53, 0.8, 0.98, 1.0)
 SKY_BOTTOM_COLOR = (0.1, 0.1, 0.43, 1.0)
@@ -369,6 +369,8 @@ def make_torus(minor_radius: float = 0.2, num_segments: int = 16) -> Geom:
 class Panda3dApp(panda3d_viewer.viewer_app.ViewerApp):
     """A Panda3D based application.
     """
+    UserCameraMask = BitMask32(2)
+
     def __init__(self,  # pylint: disable=super-init-not-called
                  config: Optional[ViewerConfig] = None) -> None:
         # Enforce viewer configuration
@@ -394,7 +396,7 @@ def __init__(self,  # pylint: disable=super-init-not-called
         config.set_value('default-directnotify-level', 'error')
         loadPrcFileData('', str(config))
 
-        # Define offscreen buffer
+        # Offscreen buffer
         self.buff: Optional[GraphicsOutput] = None
 
         # Initialize base implementation.
@@ -446,11 +448,11 @@ def keyboardInterruptHandler(
             self._make_light_ambient((0.5, 0.5, 0.5)),
             self._make_light_direct(1, (1.0, 1.0, 1.0), pos=(8.0, -8.0, 10.0))]
 
-        # Define default camera pos
-        self._camera_defaults = CAMERA_POS_DEFAULT
+        # Current camera pose
+        self._camera_defaults = CAMERA_POSE_DEFAULT
         self.reset_camera(*self._camera_defaults)
 
-        # Define clock. It will be used later to limit framerate
+        # Custom clock. It will be used later to limit framerate
         self.clock = ClockObject.get_global_clock()
         self.framerate: Optional[float] = None
 
@@ -522,26 +524,30 @@ def keyboardInterruptHandler(
         self.offA2dBottomCenter.set_pos(0, 0, self.a2dBottom)
         self.offA2dBottomRight.set_pos(self.a2dRight, 0, self.a2dBottom)
 
-        # Define widget overlay
+        # Widget overlay
         self.offscreen_graphics_lens: Optional[Lens] = None
         self.offscreen_display_region: Optional[DisplayRegion] = None
         self._help_label = None
         self._watermark: Optional[OnscreenImage] = None
         self._legend: Optional[OnscreenImage] = None
         self._clock: Optional[OnscreenText] = None
 
-        # Define input control
+        # Custom user-specified cameras
+        self._user_buffers = {}
+        self._user_cameras = {}
+
+        # Input control
         self.key_map = {"mouse1": 0, "mouse2": 0, "mouse3": 0}
 
-        # Define camera control
+        # Camera control
         self.zoom_rate = 1.03
         self.camera_lookat = np.zeros(3)
         self.longitude_deg = 0.0
         self.latitude_deg = 0.0
         self.last_mouse_x = 0.0
         self.last_mouse_y = 0.0
 
-        # Define object/highlighting selector
+        # Object/highlighting selector
         self.picker_ray: Optional[CollisionRay] = None
         self.picker_node: Optional[CollisionNode] = None
         self.picked_object: Optional[Tuple[str, str]] = None
@@ -656,7 +662,7 @@ def _open_offscreen_window(self,
         fbprops.set_rgba_bits(8, 8, 8, 8)
         fbprops.set_float_color(False)
         fbprops.set_depth_bits(16)
-        fbprops.set_float_depth(True)
+        fbprops.set_float_depth(False)
         fbprops.set_multisamples(4)
 
         # Set offscreen buffer windows properties
@@ -674,15 +680,20 @@ def _open_offscreen_window(self,
             self.pipe, "offscreen_buffer", 0, fbprops, winprops, flags,
             self.win.get_gsg(), self.win)
         if win is None:
-            raise RuntimeError("Faulty graphics pipeline of this machine.")
+            raise RuntimeError("Faulty graphics pipeline on this machine.")
         self.buff = win
 
+        # Disable automatic rendering of the buffer of efficiency
+        win.set_one_shot(True)
+
         # Append buffer to the list of windows managed by the ShowBase
         self.winList.append(win)
 
         # Attach a texture as screenshot requires copying GPU data to RAM
+        tex = Texture()
+        tex.set_format(Texture.F_rgb)
         self.buff.add_render_texture(
-            Texture(), GraphicsOutput.RTM_triggered_copy_ram)
+            tex, GraphicsOutput.RTM_triggered_copy_ram)
 
         # Create 3D camera region for the scene.
         # Set near distance of camera lens to allow seeing model from close.
@@ -742,6 +753,73 @@ def _adjust_offscreen_window_aspect_ratio(self) -> None:
         self.offA2dBottomLeft.set_pos(a2dLeft, 0, a2dBottom)
         self.offA2dBottomRight.set_pos(a2dRight, 0, a2dBottom)
 
+    def add_camera(self,
+                   name: str,
+                   size: Tuple[int, int],
+                   is_depthmap: bool) -> None:
+        # TODO: Raise exception if already exists. Add method to delete camera.
+        # TODO: Expose optional parameters to set lens type and properties.
+
+        # Create new offscreen buffer
+        fbprops = FrameBufferProperties()
+        if is_depthmap:
+            fbprops.set_depth_bits(32)
+            fbprops.set_float_depth(True)
+        else:
+            fbprops.set_rgba_bits(8, 8, 8, 8)
+            fbprops.set_float_color(False)
+            fbprops.set_depth_bits(16)
+            fbprops.set_float_depth(False)
+            fbprops.set_multisamples(4)
+        winprops = WindowProperties()
+        winprops.set_size(*size)
+        flags = GraphicsPipe.BF_refuse_window
+        buffer = self.graphicsEngine.make_output(
+            self.pipe, f"user_buffer_{name}", 0, fbprops, winprops, flags,
+            self.win.get_gsg(), self.win)
+        if buffer is None:
+            raise RuntimeError("Faulty graphics pipeline on this machine.")
+        self._user_buffers[name] = buffer
+        self.winList.append(buffer)
+
+        # Disable automatic rendering of the buffer of efficiency
+        buffer.set_one_shot(True)
+
+        # Attach a texture as screenshot requires copying GPU data to RAM
+        tex = Texture(f"user_texture_{name}")
+        if is_depthmap:
+            tex.set_format(Texture.F_depth_component)
+        else:
+            tex.set_format(Texture.F_rgb)
+        buffer.add_render_texture(
+            tex, GraphicsOutput.RTM_triggered_copy_ram)
+
+        # Create 3D camera region for the scene.
+        # See official documentation about field of view parameterization:
+        # https://docs.panda3d.org/1.10/python/programming/camera-control/perspective-lenses
+        lens = PerspectiveLens()
+        if is_depthmap:
+            lens.set_fov(50.0)   # field of view angle [0, 180], 40° by default
+            lens.set_near(0.02)  # near distance (objects closer not rendered)
+            lens.set_far(6.0)    # far distance (objects farther not rendered)
+            # lens.set_film_size(24, 36)
+            # lens.set_focal_length(50)  # Setting this will overwrite fov
+        else:
+            lens.set_near(0.1)
+        lens.set_aspect_ratio(self.get_aspect_ratio(buffer))
+        cam = self.make_camera(buffer,
+                               camName=f"user_camera_{name}",
+                               lens=lens,
+                               mask=self.UserCameraMask)
+        cam.reparent_to(self.render)
+        self._user_cameras[name] = cam
+
+        # Force rendering the scene to finalize initialization of the GSG
+        self.graphics_engine.render_frame()
+
+        # Flipped buffer upside-down
+        buffer.inverted = True
+
     def getSize(self, win: Optional[Any] = None) -> Tuple[int, int]:
         """Patched to return the size of the window used for capturing frame by
         default, instead of main window.
@@ -923,6 +1001,8 @@ def _make_axes(self) -> NodePath:
         node.set_shader_auto(True)
         node.set_light_off()
         node.hide(self.LightMask)
+        node.hide(self.UserCameraMask)
+        node.set_tag("is_virtual", "1")
         node.set_scale(0.3)
         return node
 
@@ -1074,6 +1154,8 @@ def append_frame(self,
         node.set_shader_auto(True)
         node.set_light_off()
         node.hide(self.LightMask)
+        node.hide(self.UserCameraMask)
+        node.set_tag("is_virtual", "1")
         self.append_node(root_path, name, node, frame)
 
     def append_cone(self,
@@ -1156,6 +1238,11 @@ def append_arrow(self,
         body_node.set_scale(1.0, 1.0, length)
         body_node.set_pos(0.0, 0.0, (-0.5 if anchor_top else 0.5) * length)
         arrow_node.set_scale(radius, radius, 1.0)
+
+        arrow_node.hide(self.LightMask)
+        arrow_node.hide(self.UserCameraMask)
+        arrow_node.set_tag("is_virtual", "1")
+
         self.append_node(root_path, name, arrow_node, frame)
 
     def append_mesh(self,
@@ -1575,8 +1662,11 @@ def show_node(self,
             if always_foreground is not None:
                 if always_foreground:
                     node.set_bin("fixed", 0)
+                    node.hide(self.UserCameraMask)
                 else:
                     node.clear_bin()
+                    if node.get_tag("is_virtual") != "1":
+                        node.show(self.UserCameraMask)
                 node.set_depth_test(not always_foreground)
                 node.set_depth_write(not always_foreground)
 
@@ -1593,21 +1683,27 @@ def get_camera_transform(self) -> Tuple[np.ndarray, np.ndarray]:
     def set_camera_transform(self,
                              pos: Tuple3FType,
                              quat: np.ndarray,
-                             lookat: Tuple3FType = (0.0, 0.0, 0.0)) -> None:
+                             camera_name: Optional[str] = None) -> None:
         """Set the current absolute pose of the camera.
 
         :param pos: Desired position of the camera.
         :param quat: Desired orientation of the camera as a quaternion
                      (X, Y, Z, W).
-        :param lookat: Point at which the camera is looking at. It is partially
-                       redundant with the desired orientation and will take
-                       precedence in case of inconsistency. It is also involved
-                       in zoom control.
+        :param camera_name: TODO: Write documentation.
         """
-        self.camera.set_pos(*pos)
-        self.camera.set_quat(LQuaternion(quat[-1], *quat[:-1]))
-        self.camera_lookat = np.array(lookat)
-        self.move_orbital_camera_task()
+        # Pick the right camera
+        if camera_name is None:
+            camera = self.camera
+        else:
+            camera = self._user_cameras[camera_name]
+
+        # Move the camera
+        camera.set_pos_quat(Vec3(*pos), Quat(quat[-1], *quat[:-1]))
+
+        # Reset orbital camera control
+        if camera_name is None:
+            self.camera_lookat = np.array([0.0, 0.0, 0.0])
+            self.move_orbital_camera_task()
 
     def get_camera_lookat(self) -> np.ndarray:
         """Get the location of the point toward which the camera is looking at.
@@ -1674,6 +1770,7 @@ def save_screenshot(self, filename: Optional[str] = None) -> bool:
         # It will take into account the updated position of the camera.
         assert self.buff is not None
         self.buff.trigger_copy()
+        self.buff.set_one_shot(True)
         self.graphics_engine.render_frame()
 
         # Capture frame as image
@@ -1696,8 +1793,8 @@ def save_screenshot(self, filename: Optional[str] = None) -> bool:
         return True
 
     def get_screenshot(self,
-                       requested_format: str = 'RGB',
-                       raw: bool = False) -> Union[np.ndarray, bytes]:
+                       camera_name: Optional[str] = None
+                       ) -> Union[np.ndarray, bytes]:
         """Patched to take screenshot of the last window available instead of
         the main one, and to add raw data return mode for efficient
         multiprocessing.
@@ -1708,32 +1805,34 @@ def get_screenshot(self,
             scheduler. The framerate limit must be disable manually to avoid
             such limitation.
 
-        .. note::
-            Internally, Panda3d uses BGRA, so using it is slightly faster than
-            RGBA, but not RGB since there is one channel missing.
-
-        :param requested_format: Desired export format (e.g. 'RGB' or 'BGRA')
-        :param raw: whether to return a raw memory view of bytes, of a
-                    structured `np.ndarray` of uint8 with dimensions [H, W, D].
+        :param camera_name: TODO: Write documentaytion.
         """
+        # Get desired buffer
+        if camera_name is None:
+            assert self.buff is not None
+            buffer = self.buff
+        else:
+            buffer = self._user_buffers[camera_name]
+
         # Refresh the scene
-        assert self.buff is not None
-        self.buff.trigger_copy()
+        buffer.trigger_copy()
+        buffer.set_one_shot(True)
         self.graphics_engine.render_frame()
 
         # Get frame as raw texture
-        assert self.buff is not None
-        texture = self.buff.get_texture()
+        texture = buffer.get_texture()
 
         # Extract raw array buffer from texture
-        image = texture.get_ram_image_as(requested_format)
-
-        # Return raw buffer if requested
-        if raw:
-            return image.get_data()
+        is_depth_map = texture.format == Texture.F_depth_component32
+        if is_depth_map:
+            image = texture.get_ram_image()
+        else:
+            image = texture.get_ram_image_as('RGB')
 
         # Convert raw texture to numpy array if requested
         xsize, ysize = texture.get_x_size(), texture.get_y_size()
+        if is_depth_map:
+            return np.frombuffer(image, np.float32).reshape((ysize, xsize))
         return np.frombuffer(image, np.uint8).reshape((ysize, xsize, -1))
 
     def enable_shadow(self, enable: bool) -> None:

python/jiminy_py/src/jiminy_py/viewer/replay.py

@@ -595,6 +595,7 @@ def play_trajectories(
 
             # Create frame storage
             frame = av.VideoFrame(*record_video_size, 'rgb24')
+            frame_bytes = memoryview(frame.planes[0])
 
         # Add frames to video sequentially
         update_hook_t = None
@@ -636,9 +637,8 @@ def play_trajectories(
                     # Update frame.
                     # Note that `capture_frame` is by far the main bottleneck
                     # of the whole recording process (~75% on discrete gpu).
-                    buffer = Viewer.capture_frame(
-                        *record_video_size, raw_data=True)
-                    memoryview(frame.planes[0])[:] = buffer
+                    buffer = Viewer.capture_frame(*record_video_size)
+                    frame_bytes[:] = buffer.reshape(-1).data
 
                     # Write frame
                     for packet in stream.encode(frame):