[FEATURE] Support attaching camera sensors with Rasterizer batched rendering. (#2212)

Author: Narsil

genesis/engine/sensors/camera.py

@@ -267,11 +267,6 @@ def move_to_attach(self):
         link_pos = self._link.get_pos()
         link_quat = self._link.get_quat()
 
-        # Handle batched case - use first environment
-        if link_pos.ndim > 1:
-            link_pos = link_pos[0]
-            link_quat = link_quat[0]
-
         link_T = trans_quat_to_T(link_pos, link_quat)
         camera_T = torch.matmul(link_T, offset_T)
 
@@ -448,9 +443,6 @@ def _create_standalone_context(self, scene):
                 )
             env_separate_rigid = False
         else:
-            if self._link is not None:
-                gs.raise_exception("Rasterizer with n_envs > 1, does not work with attached cameras yet.")
-
             if scene.n_envs > 1:
                 gs.logger.warning(
                     "Rasterizer with n_envs > 1 is slow as it doesn't do batched rendering consider using BatchRenderer instead."
@@ -498,11 +490,6 @@ def _update_camera_pose(self):
             link_pos = self._link.get_pos()
             link_quat = self._link.get_quat()
 
-            # Handle batched case - use first environment
-            if link_pos.ndim > 1:
-                link_pos = link_pos[0]
-                link_quat = link_quat[0]
-
             # Apply pos directly as offset from link
             from genesis.utils.geom import transform_by_quat
 
@@ -609,11 +596,6 @@ def build(self):
             link_pos = self._link.get_pos()
             link_quat = self._link.get_quat()
 
-            # Handle batched case - use first environment
-            if link_pos.ndim > 1:
-                link_pos = link_pos[0]
-                link_quat = link_quat[0]
-
             # Apply pos directly as offset from link
             from genesis.utils.geom import transform_by_quat
 
@@ -758,8 +740,7 @@ def build(self):
             resolutions = [s._options.res for s in all_sensors]
             if len(set(resolutions)) > 1:
                 gs.raise_exception(
-                    f"All BatchRendererCameraSensor instances must have the same resolution. "
-                    f"Found: {set(resolutions)}"
+                    f"All BatchRendererCameraSensor instances must have the same resolution. Found: {set(resolutions)}"
                 )
 
             br_options = BatchRendererOptions(

genesis/ext/pyrender/renderer.py

@@ -354,9 +354,9 @@ def _floor_pass(self, scene, flags, seg_node_map=None, env_idx=-1):
         glClearColor(0.0, 0.0, 0.0, 1.0)
         glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 
-        V, P = self._get_camera_matrices(scene)
+        V, P = self._get_camera_matrices(scene, env_idx)
+        cam_pose = self._get_camera_pose(scene, env_idx)[:3, 3]
 
-        cam_pos = scene.get_pose(scene.main_camera_node)[:3, 3]
         screen_size = np.array([self.viewport_width, self.viewport_height], np.float32)
 
         self.jit.forward_pass(
@@ -391,8 +391,8 @@ def _forward_pass(self, scene, flags, seg_node_map=None, env_idx=-1):
             glEnable(GL_MULTISAMPLE)
 
         # Set up camera matrices
-        V, P = self._get_camera_matrices(scene)
-        cam_pos = scene.get_pose(scene.main_camera_node)[:3, 3]
+        V, P = self._get_camera_matrices(scene, env_idx)
+        cam_pos = self._get_camera_pose(scene, env_idx)[:3, 3]
 
         floor_tex = self._floor_texture_color._texid if flags & RenderFlags.REFLECTIVE_FLOOR else 0
         screen_size = np.array([self.viewport_width, self.viewport_height], np.float32)
@@ -455,7 +455,7 @@ def _normal_pass(self, scene, flags, env_idx=-1):
         program = None
 
         # Set up camera matrices
-        V, P = self._get_camera_matrices(scene)
+        V, P = self._get_camera_matrices(scene, env_idx)
 
         # Now, render each object in sorted order
         for node in scene.sorted_mesh_nodes():
@@ -690,15 +690,25 @@ def _reset_active_textures(self):
     # Camera Matrix Management
     ###########################################################################
 
-    def _get_camera_matrices(self, scene):
+    def _get_camera_matrices(self, scene, env_idx):
         main_camera_node = scene.main_camera_node
         if main_camera_node is None:
             raise ValueError("Cannot render scene without a camera")
         P = main_camera_node.camera.get_projection_matrix(width=self.viewport_width, height=self.viewport_height)
-        pose = scene.get_pose(main_camera_node)
+        pose = self._get_camera_pose(scene, env_idx)
         V = np.linalg.inv(pose)  # V maps from world to camera
         return V, P
 
+    def _get_camera_pose(self, scene, env_idx):
+        cam_pos = scene.get_pose(scene.main_camera_node)
+        if len(cam_pos.shape) == 3:
+            if cam_pos.shape[0] != 1:
+                assert env_idx != -1, "We have a multiple camera pose scene, we should be rendering per env"
+                cam_pos = cam_pos[env_idx]
+            else:
+                cam_pos = cam_pos[0]
+        return cam_pos
+
     def _get_light_cam_matrices(self, scene, light_node, flags):
         light = light_node.light
         pose = scene.get_pose(light_node)

genesis/ext/pyrender/scene.py

@@ -89,7 +89,7 @@ def __init__(self, nodes=None, bg_color=None, ambient_light=None, n_envs=None, n
             for node in nodes:
                 for child in node.children:
                     if node_parent_map[child] is not None:
-                        raise ValueError("Nodes may not have more than " "one parent")
+                        raise ValueError("Nodes may not have more than one parent")
                     node_parent_map[child] = node
             for node in node_parent_map:
                 if node_parent_map[node] is None:
@@ -294,7 +294,7 @@ def add(self, obj, name=None, pose=None, parent_node=None, parent_name=None):
 
         if parent_node is None and parent_name is not None:
             try:
-                parent_node, = self.get_nodes(name=parent_name)
+                (parent_node,) = self.get_nodes(name=parent_name)
             except ValueError:
                 if len(parent_nodes) == 0:
                     raise ValueError(f"No parent node with name '{parent_name}' found")
@@ -601,7 +601,9 @@ def from_trimesh_scene(trimesh_scene, bg_color=None, ambient_light=None):
         return scene_pr
 
     def sorted_mesh_nodes(self):
-        cam_loc = self.get_pose(self.main_camera_node)[:3, 3]
+        cam_pos = self.get_pose(self.main_camera_node)
+        cam_loc = cam_pos[..., :3, 3]
+        batched_pos = len(cam_pos.shape) == 3
         solid_nodes = []
         trans_nodes = []
         for node in self.mesh_nodes:
@@ -612,7 +614,12 @@ def sorted_mesh_nodes(self):
                 solid_nodes.append(node)
 
         # TODO BETTER SORTING METHOD
-        trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
-        solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
+        if batched_pos:
+            # FIXME normally sorting should be done PER scene when having a batched rasterizer render
+            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
+            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
+        else:
+            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
+            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
 
         return solid_nodes + trans_nodes

pyproject.toml

@@ -139,6 +139,10 @@ force-exclude = '''
 )/
 '''
 
+
+[tool.ruff]
+line-length = 120
+
 [tool.pytest.ini_options]
 addopts = [
     "--color=yes",

tests/test_sensor_camera.py

@@ -160,7 +160,6 @@ def test_rasterizer_camera_sensor_n_envs(show_viewer, png_snapshot):
     # Add a plane
     scene.add_entity(
         morph=gs.morphs.Plane(),
-        surface=gs.surfaces.Rough(color=(0.4, 0.4, 0.4)),
     )
 
     # Add a sphere
@@ -193,9 +192,15 @@ def test_rasterizer_camera_sensor_n_envs(show_viewer, png_snapshot):
 
 @pytest.mark.required
 @pytest.mark.skipif(sys.platform == "darwin", reason="Not supported on this machine because it requires OpenGL 4.2.")
-def test_rasterizer_camera_sensor_n_envs_attached_camera():
-    scene = gs.Scene()
+def test_rasterizer_camera_sensor_n_attached_camera(show_viewer, png_snapshot):
+    scene = gs.Scene(show_viewer=show_viewer)
 
+    # Add a plane
+    scene.add_entity(
+        morph=gs.morphs.Plane(),
+    )
+
+    # Add a sphere
     sphere = scene.add_entity(
         morph=gs.morphs.Sphere(
             radius=0.3,
@@ -206,11 +211,29 @@ def test_rasterizer_camera_sensor_n_envs_attached_camera():
         ),
     )
 
-    scene.add_sensor(
-        gs.sensors.RasterizerCameraOptions(
-            entity_idx=sphere.idx,
-        )
+    options = gs.sensors.RasterizerCameraOptions(
+        res=(64, 64),
+        pos=(-0.4, 0.1, 2.0),
+        lookat=(-0.6, 0.4, 1.0),
+        fov=60.0,
+        entity_idx=sphere.idx,
+        draw_debug=show_viewer,
     )
+    camera = scene.add_sensor(options)
+
+    scene.build(n_envs=2)
 
-    with pytest.raises(gs.GenesisException, match="does not work with attached cameras yet."):
-        scene.build(n_envs=2)
+    # Disable shadows systematically for Rasterizer because they are forcibly disabled on CPU backend anyway
+    camera._shared_metadata.context.shadow = False
+
+    sphere.set_pos([[0.0, 0.0, 1.0], [0.2, 0.0, 0.5]])
+    scene.step()
+
+    data = camera.read()
+
+    assert data.rgb.shape == (2, 64, 64, 3)
+    assert data.rgb.dtype == torch.uint8
+    assert (data.rgb[0] != data.rgb[1]).any(), "We should have different frames"
+
+    for i in range(scene.n_envs):
+        assert rgb_array_to_png_bytes(data.rgb[i]) == png_snapshot

tests/utils.py

@@ -36,7 +36,7 @@
 DEFAULT_BRANCH_NAME = "main"
 
 HUGGINGFACE_ASSETS_REVISION = "701f78c1465f0a98f6540bae6c9daacaa551b7bf"
-HUGGINGFACE_SNAPSHOT_REVISION = "1df3aa3732abcf9b1701e180e2d266ac9d8d411e"
+HUGGINGFACE_SNAPSHOT_REVISION = "f13f28423a8961072832ead74df5d5703e01923e"
 
 MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS)
 IMAGE_EXTENSIONS = (".png", ".jpg")