Fix point cloud generation from depth map.

Author: duburcqa

genesis/vis/camera.py

@@ -471,21 +471,23 @@ def get_segmentation_idx_dict(self):
     @gs.assert_built
     def render_pointcloud(self, world_frame=True):
         """
-        Render a partial point cloud from the camera view. Returns a (res[0], res[1], 3) numpy array representing the point cloud in each pixel.
+        Render a partial point cloud from the camera view.
+
         Parameters
         ----------
         world_frame : bool, optional
             Whether the point cloud is on camera frame or world frame.
+
         Returns
         -------
         pc : np.ndarray
-            the point cloud
+            Numpy array of shape (res[0], res[1], 3) representing the point cloud in each pixel.
         mask_arr : np.ndarray
             The valid depth mask.
         """
         self._rasterizer.update_scene()
         rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._rasterizer.render_camera(
-            self, False, True, False, normal=False
+            self, rgb=False, depth=True, segmentation=False, normal=False
         )
 
         def opengl_projection_matrix_to_intrinsics(P: np.ndarray, width: int, height: int):
@@ -536,16 +538,11 @@ def backproject_depth_to_pointcloud(K: np.ndarray, depth: np.ndarray, pose, worl
             world_y = normalized_y * depth_grid / _fy
             world_z = depth_grid
 
-            pc = np.stack((world_x, world_y, world_z), axis=1)
-
-            point_cloud_h = np.concatenate((pc, np.ones((len(pc), 1), dtype=np.float32)), axis=1)
+            point_cloud = np.stack((world_x, world_y, world_z, np.ones((depth.size,), dtype=np.float32)), axis=-1)
             if world:
-                point_cloud_world = point_cloud_h @ pose.T
-                point_cloud_world = point_cloud_world[:, :3].reshape((*depth.shape, 3))
-                return point_cloud_world, mask
-            else:
-                point_cloud = point_cloud_h[:, :3].reshape((*depth.shape, 3))
-                return point_cloud, mask
+                point_cloud = point_cloud @ pose.T
+            point_cloud = point_cloud[:, :3].reshape((*depth.shape, 3))
+            return point_cloud, mask
 
         intrinsic_K = opengl_projection_matrix_to_intrinsics(
             self._rasterizer._camera_nodes[self.uid].camera.get_projection_matrix(),

tests/test_render.py

@@ -271,6 +271,77 @@ def test_batched_offscreen_rendering(show_viewer, tol):
             assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol)
 
 
+@pytest.mark.required
+def test_point_cloud(show_viewer):
+    CAMERA_DIST = 8.0
+    OBJ_OFFSET = 10.0
+    BOX_HALFSIZE = 1.0
+    SPHERE_RADIUS = 1.0
+
+    scene = gs.Scene(
+        show_viewer=show_viewer,
+        show_FPS=False,
+    )
+    scene.add_entity(
+        morph=gs.morphs.Sphere(
+            pos=(0.0, OBJ_OFFSET, 0.0),
+            radius=SPHERE_RADIUS,
+            fixed=True,
+        ),
+    )
+    camera_sphere = scene.add_camera(
+        pos=(0.0, OBJ_OFFSET, CAMERA_DIST),
+        lookat=(0.0, OBJ_OFFSET, 0.0),
+        GUI=show_viewer,
+    )
+    scene.add_entity(
+        morph=gs.morphs.Box(
+            pos=(0.0, -OBJ_OFFSET, 0.0),
+            size=(2.0 * BOX_HALFSIZE, 2.0 * BOX_HALFSIZE, 2.0 * BOX_HALFSIZE),
+            fixed=True,
+        )
+    )
+    camera_box_1 = scene.add_camera(
+        pos=(0.0, -OBJ_OFFSET, CAMERA_DIST),
+        lookat=(0.0, -OBJ_OFFSET, 0.0),
+        GUI=show_viewer,
+    )
+    camera_box_2 = scene.add_camera(
+        pos=np.array((CAMERA_DIST, CAMERA_DIST - OBJ_OFFSET, CAMERA_DIST)),
+        lookat=(0.0, -OBJ_OFFSET, 0.0),
+        GUI=show_viewer,
+    )
+    for camera in scene.visualizer.cameras:
+        camera._near = 1.0
+    scene.build()
+
+    if show_viewer:
+        for camera in scene.visualizer.cameras:
+            camera.render(rgb=True, depth=True)
+
+    point_cloud, mask = camera_box_1.render_pointcloud(world_frame=False)
+    point_cloud = point_cloud[mask]
+    assert_allclose(CAMERA_DIST - point_cloud[:, 2], BOX_HALFSIZE, atol=1e-4)
+    assert np.all(-BOX_HALFSIZE <= point_cloud[:, :2].min(axis=0))
+    assert np.all(point_cloud[:, :2].max(axis=0) <= BOX_HALFSIZE)
+
+    point_cloud, mask = camera_box_2.render_pointcloud(world_frame=False)
+    point_cloud = point_cloud[mask]
+    point_cloud = point_cloud @ gu.z_up_to_R(np.array((1.0, 1.0, 1.0)), np.array((0.0, 0.0, 1.0))).T
+    point_cloud -= np.array((CAMERA_DIST, CAMERA_DIST, CAMERA_DIST))
+    assert_allclose(np.linalg.norm(point_cloud, ord=float("inf"), axis=-1), BOX_HALFSIZE, atol=1e-4)
+
+    point_cloud, mask = camera_box_2.render_pointcloud(world_frame=True)
+    point_cloud = point_cloud[mask]
+    point_cloud += np.array((0.0, OBJ_OFFSET, 0.0))
+    assert_allclose(np.linalg.norm(point_cloud, ord=float("inf"), axis=-1), BOX_HALFSIZE, atol=1e-4)
+
+    # It is not possible to get higher accuracy because of tesselation
+    point_cloud, mask = camera_sphere.render_pointcloud(world_frame=False)
+    point_cloud = point_cloud[mask]
+    assert_allclose(np.linalg.norm((0.0, 0.0, CAMERA_DIST) - point_cloud, axis=-1), SPHERE_RADIUS, atol=1e-2)
+
+
 @pytest.mark.required
 def test_batched_mounted_camera_rendering(show_viewer, tol):
     scene = gs.Scene(