[BUG FIX] Fix point cloud rendering from depth map (cont'd). (Genesis-Embodied-AI#1515)

* Fix point cloud generation from depth map.
* Expose depth-only rendering option, w/o MSAA for invertible mapping.
* Fix camera render API.

Author: duburcqa

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -106,7 +106,7 @@ def __init__(
         if len(self._sdf_faces) > 50000:
             mesh_descr = f"({mesh.metadata['mesh_path']})" if "mesh_path" in mesh.metadata else ""
             gs.logger.warning(
-                "Beware that SDF pre-processing of mesh {mesh_descr} having more than 50000 vertices may take a very "
+                f"Beware that SDF pre-processing of mesh {mesh_descr} having more than 50000 vertices may take a very "
                 "long time (>10min) and require large RAM allocation (>20Gb). Please either enable convexify or "
                 "decimation. (see FileMorph options)"
             )

genesis/ext/pyrender/camera.py

@@ -109,11 +109,7 @@ class PerspectiveCamera(Camera):
     """
 
     def __init__(self, yfov, znear=DEFAULT_Z_NEAR, zfar=None, aspectRatio=None, name=None):
-        super(PerspectiveCamera, self).__init__(
-            znear=znear,
-            zfar=zfar,
-            name=name,
-        )
+        super(PerspectiveCamera, self).__init__(znear=znear, zfar=zfar, name=name)
 
         self.yfov = yfov
         self.aspectRatio = aspectRatio
@@ -213,11 +209,7 @@ class OrthographicCamera(Camera):
     """
 
     def __init__(self, xmag, ymag, znear=DEFAULT_Z_NEAR, zfar=DEFAULT_Z_FAR, name=None):
-        super(OrthographicCamera, self).__init__(
-            znear=znear,
-            zfar=zfar,
-            name=name,
-        )
+        super(OrthographicCamera, self).__init__(znear=znear, zfar=zfar, name=name)
 
         self.xmag = xmag
         self.ymag = ymag
@@ -313,11 +305,7 @@ class IntrinsicsCamera(Camera):
     """
 
     def __init__(self, fx, fy, cx, cy, znear=DEFAULT_Z_NEAR, zfar=DEFAULT_Z_FAR, name=None):
-        super(IntrinsicsCamera, self).__init__(
-            znear=znear,
-            zfar=zfar,
-            name=name,
-        )
+        super(IntrinsicsCamera, self).__init__(znear=znear, zfar=zfar, name=name)
 
         self.fx = fx
         self.fy = fy

genesis/ext/pyrender/offscreen.py

@@ -104,14 +104,15 @@ def render(
         self,
         scene,
         renderer,
+        rgb=True,
+        seg=False,
+        normal=False,
+        depth=False,
         flags=RenderFlags.NONE,
         camera_node=None,
         shadow=False,
-        seg=False,
-        ret_depth=False,
         plane_reflection=False,
         env_separate_rigid=False,
-        normal=False,
     ):
         """Render a scene with the given set of flags.
 
@@ -131,6 +132,9 @@ def render(
         depth_im : (h, w) float32
             The depth buffer in linear units.
         """
+        if seg and rgb:
+            gs.raise_exception("RGB and segmentation map cannot be rendered in the same forward pass.")
+
         if not self._has_valid_context:
             gs.raise_exception(
                 "Ensure that the right context is set before rendering. Please call the method 'make_current'."
@@ -144,6 +148,9 @@ def render(
         if shadow and not self._is_software:
             flags |= RenderFlags.SHADOWS_ALL
 
+        if depth and not (rgb or seg):
+            flags |= RenderFlags.DEPTH_ONLY
+
         if plane_reflection and not self._is_software:
             flags |= RenderFlags.REFLECTIVE_FLOOR
 
@@ -156,23 +163,30 @@ def render(
         else:
             seg_node_map = None
 
-        if ret_depth:
+        if depth:
             flags |= RenderFlags.RET_DEPTH
 
-        if self._platform.supports_framebuffers():
-            flags |= RenderFlags.OFFSCREEN
-            retval = renderer.render(scene, flags, seg_node_map)
-        else:
-            renderer.render(scene, flags, seg_node_map)
-            depth = renderer.read_depth_buf()
-            if flags & RenderFlags.DEPTH_ONLY:
-                retval = depth
+        if rgb or depth or seg:
+            if self._platform.supports_framebuffers():
+                flags |= RenderFlags.OFFSCREEN
+                retval = renderer.render(scene, flags, seg_node_map)
             else:
-                color = renderer.read_color_buf()
-                retval = color, depth
+                if flags & RenderFlags.ENV_SEPARATE:
+                    gs.raise_exception("'env_separate_rigid=True' not supported on this platform.")
+                if normal:
+                    gs.raise_exception("'normal=True' not supported on this platform.")
+                renderer.render(scene, flags, seg_node_map)
+                if depth:
+                    depth = renderer.read_depth_buf()
+                if flags & RenderFlags.DEPTH_ONLY:
+                    retval = (depth,)
+                else:
+                    color = renderer.read_color_buf()
+                    retval = (color, depth)
+        else:
+            retval = ()
 
         if normal:
-
             class CustomShaderCache:
                 def __init__(self):
                     self.program = None
@@ -193,8 +207,8 @@ def get_program(self, vertex_shader, fragment_shader, geometry_shader=None, defi
             if env_separate_rigid:
                 flags |= RenderFlags.ENV_SEPARATE
             glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
-            normal_arr, _ = renderer.render(scene, flags, is_first_pass=False)
-            retval = retval + (normal_arr,)
+            normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
+            retval = (*retval, normal_arr)
 
             renderer._program_cache = old_cache
 

genesis/ext/pyrender/renderer.py

@@ -115,7 +115,7 @@ def point_size(self):
     def point_size(self, value):
         self._point_size = float(value)
 
-    def render(self, scene, flags, seg_node_map=None, is_first_pass=True):
+    def render(self, scene, flags, seg_node_map=None, *, is_first_pass=True, force_skip_shadows=False):
         """Render a scene with the given set of flags.
 
         Parameters
@@ -144,7 +144,7 @@ def render(self, scene, flags, seg_node_map=None, is_first_pass=True):
             self._update_context(scene, flags)
             self.jit.update(scene)
 
-        if flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.SEG or flags & RenderFlags.FLAT:
+        if flags & RenderFlags.SEG or flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.FLAT:
             flags &= ~RenderFlags.REFLECTIVE_FLOOR
 
         if flags & RenderFlags.ENV_SEPARATE and flags & RenderFlags.OFFSCREEN:
@@ -159,7 +159,7 @@ def render(self, scene, flags, seg_node_map=None, is_first_pass=True):
             env_idx = i if use_env_idx else -1
 
             # Render necessary shadow maps
-            if not (flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.SEG):
+            if not (force_skip_shadows or flags & RenderFlags.SEG or flags & RenderFlags.DEPTH_ONLY):
                 for ln in scene.light_nodes:
                     take_pass = False
                     if isinstance(ln.light, DirectionalLight) and flags & RenderFlags.SHADOWS_DIRECTIONAL:
@@ -180,37 +180,28 @@ def render(self, scene, flags, seg_node_map=None, is_first_pass=True):
                 self._floor_pass(scene, flags, env_idx=env_idx)
 
             retval = self._forward_pass(scene, flags, seg_node_map=seg_node_map, env_idx=env_idx)
-            if isinstance(retval, tuple):
+            if retval is not None:
                 if retval_list is None:
-                    retval_list = tuple([[val] for val in retval])
+                    retval_list = tuple([val] for val in retval)
                 else:
                     for idx, val in enumerate(retval):
                         retval_list[idx].append(val)
-            elif retval is not None:
-                if retval_list is None:
-                    retval_list = [retval]
-                else:
-                    retval_list.append(retval)
 
             # If necessary, make normals pass
             if flags & (RenderFlags.VERTEX_NORMALS | RenderFlags.FACE_NORMALS):
-                self._normals_pass(scene, flags, env_idx=env_idx)
-
-        if use_env_idx:
-            if isinstance(retval_list, list):
-                retval_list = np.stack(retval_list, axis=0)
-            elif isinstance(retval_list, tuple):
-                retval_list = tuple([np.stack(val_list, axis=0) for val_list in retval_list])
-        else:
-            if isinstance(retval_list, list):
-                retval_list = retval_list[0]
-            elif isinstance(retval_list, tuple):
-                retval_list = tuple([val_list[0] for val_list in retval_list])
+                self._normal_pass(scene, flags, env_idx=env_idx)
 
         # Update camera settings for retrieving depth buffers
         self._latest_znear = scene.main_camera_node.camera.znear
         self._latest_zfar = scene.main_camera_node.camera.zfar
 
+        if retval_list is None:
+            return
+
+        if use_env_idx:
+            retval_list = tuple(np.stack(val_list, axis=0) for val_list in retval_list)
+        else:
+            retval_list = tuple([val_list[0] for val_list in retval_list])
         return retval_list
 
     def render_text(
@@ -313,114 +304,6 @@ def render_texts(self, texts, x, y, font_name="UbuntuMono-Regular", font_pt=40,
         for i, text in enumerate(texts):
             font.render_string(text, x, int(y - i * font_pt * 1.1), scale, TextAlign.TOP_LEFT)
 
-    def read_depth_buf(self):
-        """Read and return the current viewport's color buffer.
-
-        Returns
-        -------
-        depth_im : (h, w) float32
-            The depth buffer in linear units.
-        """
-        width, height = self.viewport_width, self.viewport_height
-        glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
-        glReadBuffer(GL_FRONT)
-        depth_buf = glReadPixels(
-            0, 0, width, height, GL_DEPTH_COMPONENT, GL_FLOAT
-        )
-
-        depth_im = np.frombuffer(depth_buf, dtype=np.float32)
-        depth_im = depth_im.reshape((height, width))
-        depth_im = np.flip(depth_im, axis=0)
-
-        inf_inds = (depth_im == 1.0)
-        depth_im = 2.0 * depth_im - 1.0
-        z_near, z_far = self._latest_znear, self._latest_zfar
-        noninf = np.logical_not(inf_inds)
-        if z_far is None:
-            depth_im[noninf] = 2 * z_near / (1.0 - depth_im[noninf])
-        else:
-            depth_im[noninf] = ((2.0 * z_near * z_far) /
-                                (z_far + z_near - depth_im[noninf] *
-                                (z_far - z_near)))
-        depth_im[inf_inds] = 0.0
-
-        # Resize for macos if needed
-        if sys.platform == 'darwin':
-            depth_im = self._resize_image(depth_im)
-
-        return depth_im
-
-    def read_color_buf(self):
-        """Read and return the current viewport's color buffer.
-
-        Alpha cannot be computed for an on-screen buffer.
-
-        Returns
-        -------
-        color_im : (h, w, 3) uint8
-            The color buffer in RGB byte format.
-        """
-        # Extract color image from frame buffer
-        width, height = self.viewport_width, self.viewport_height
-        glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
-        glReadBuffer(GL_FRONT)
-        color_buf = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE)
-
-        # Re-format them into numpy arrays
-        color_im = np.frombuffer(color_buf, dtype=np.uint8)
-        color_im = color_im.reshape((height, width, 3))
-        color_im = np.flip(color_im, axis=0)
-
-        # Resize for macos if needed
-        if sys.platform == "darwin":
-            color_im = self._resize_image(color_im, True)
-
-        return color_im
-
-    def get_tex_image(self, tex_id, bind_target=GL_TEXTURE_2D, read_target=GL_TEXTURE_2D, width=None, height=None):
-        """Read and return the current viewport's color buffer.
-
-        Returns
-        -------
-        depth_im : (h, w) float32
-            The depth buffer in linear units.
-        """
-        # width, height = self.viewport_width, self.viewport_height
-        # width, height = 8192, 8192
-        # glBindFramebuffer(GL_READ_FRAMEBUFFER, self._shadow_fb)
-        # glReadBuffer(GL_NONE)
-        # depth_buf = glReadPixels(
-        #     0, 0, width, height, GL_DEPTH_COMPONENT, GL_FLOAT
-        # )
-
-        glBindTexture(bind_target, tex_id)
-        depth_buf = glGetTexImage(read_target, 0, GL_RGB, GL_FLOAT)
-
-        depth_im = np.frombuffer(depth_buf, dtype=np.float32)
-        if width is None:
-            height = int(depth_im.size**0.5)
-            width = height
-        depth_im = depth_im.reshape((height, width, 3))
-        depth_im = np.flip(depth_im, axis=0)
-
-        return depth_im
-
-        inf_inds = depth_im == 1.0
-        depth_im = 2.0 * depth_im - 1.0
-        z_near, z_far = self._latest_znear, self._latest_zfar
-        noninf = np.logical_not(inf_inds)
-        if z_far is None:
-            depth_im[noninf] = 2 * z_near / (1.0 - depth_im[noninf])
-        else:
-            depth_im[noninf] = (2.0 * z_near * z_far) / (z_far + z_near - depth_im[noninf] * (z_far - z_near))
-        depth_im[inf_inds] = 0.0
-
-        # # Resize for macos if needed
-        # if sys.platform == 'darwin':
-        #     depth_im = self._resize_image(depth_im)
-
-        return depth_im
-
     def delete(self):
         """Free all allocated OpenGL resources."""
         # Free shaders
@@ -505,7 +388,7 @@ def _forward_pass(self, scene, flags, seg_node_map=None, env_idx=-1):
 
         glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
-        if flags & RenderFlags.SEG:
+        if flags & RenderFlags.SEG or flags & RenderFlags.DEPTH_ONLY:
             glDisable(GL_MULTISAMPLE)
         else:
             glEnable(GL_MULTISAMPLE)
@@ -569,7 +452,7 @@ def _shadow_mapping_pass(self, scene, light_node, flags, env_idx=-1):
 
         self.jit.shadow_mapping_pass(self, V, P, flags, ProgramFlags.NONE, env_idx=env_idx)
 
-    def _normals_pass(self, scene, flags, env_idx=-1):
+    def _normal_pass(self, scene, flags, env_idx=-1):
         # Set up viewport for render
         self._configure_forward_pass_viewport(flags)
         program = None
@@ -990,7 +873,7 @@ def _configure_forward_pass_viewport(self, flags):
         # If using offscreen render, bind main framebuffer
         if flags & RenderFlags.OFFSCREEN:
             self._configure_main_framebuffer()
-            if flags & RenderFlags.SEG:
+            if flags & RenderFlags.SEG or flags & RenderFlags.DEPTH_ONLY:
                 glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb)
             else:
                 glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb_ms)
@@ -1149,7 +1032,7 @@ def _delete_main_framebuffer(self):
     def _read_main_framebuffer(self, scene, flags):
         width, height = self._main_fb_dims
 
-        if not flags & RenderFlags.SEG:
+        if not (flags & RenderFlags.SEG or flags & RenderFlags.DEPTH_ONLY):
             # Bind framebuffer and blit buffers
             glBindFramebuffer(GL_READ_FRAMEBUFFER, self._main_fb_ms)
             glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb)
@@ -1167,19 +1050,19 @@ def _read_main_framebuffer(self, scene, flags):
 
             # Resize
             depth_im = self._resize_image(depth_im, antialias=False)
-        else:
-            depth_im = None
 
         if flags & RenderFlags.DEPTH_ONLY:
-            return depth_im
+            return (depth_im,)
 
         # Read color
         color_im = self.jit.read_color_buf(width, height, flags & RenderFlags.RGBA)
 
         # Resize
         color_im = self._resize_image(color_im, antialias=not flags & RenderFlags.SEG)
 
-        return color_im, depth_im
+        if flags & RenderFlags.RET_DEPTH:
+            return color_im, depth_im
+        return (color_im,)
 
     def _resize_image(self, value, antialias):
         """Rescale the generated image if necessary."""

genesis/ext/pyrender/texture.py

@@ -163,6 +163,7 @@ def delete(self):
     ##################
     # OpenGL code
     ##################
+
     def _add_to_context(self):
         if self._texid is not None:
             return