diff --git a/genesis/engine/sensors/camera.py b/genesis/engine/sensors/camera.py index 3252331e01..16d4814892 100644 --- a/genesis/engine/sensors/camera.py +++ b/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( diff --git a/genesis/ext/pyrender/renderer.py b/genesis/ext/pyrender/renderer.py index a3d7c2ed1a..50a44aed94 100644 --- a/genesis/ext/pyrender/renderer.py +++ b/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) diff --git a/genesis/ext/pyrender/scene.py b/genesis/ext/pyrender/scene.py index 0c1d788e42..5c8e5fac0a 100644 --- a/genesis/ext/pyrender/scene.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 953e1f519e..a0dc25dd09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,10 @@ force-exclude = ''' )/ ''' + +[tool.ruff] +line-length = 120 + [tool.pytest.ini_options] addopts = [ "--color=yes", diff --git a/tests/test_sensor_camera.py b/tests/test_sensor_camera.py index 5a57a4c847..ecba31330f 100644 --- a/tests/test_sensor_camera.py +++ b/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 diff --git a/tests/utils.py b/tests/utils.py index e6f1226928..cbf1dfc1c9 100644 --- a/tests/utils.py +++ b/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")