diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 57aa1afdff..20ba6016cd 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -1,5 +1,6 @@ import os import pickle +import sys import time import numpy as np @@ -526,7 +527,7 @@ def add_camera( focus_dist=None, GUI=False, spp=256, - denoise=True, + denoise=None, env_idx=None, ): """ @@ -564,14 +565,16 @@ def add_camera( Samples per pixel. Only available when using RayTracer renderer. Defaults to 256. denoise : bool Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults - to True. If OptiX denoiser is not available in your platform, consider enabling the OIDN denoiser option - when building the RayTracer. + to True on Linux, otherwise False. If OptiX denoiser is not available in your platform, consider enabling + the OIDN denoiser option when building the RayTracer. Returns ------- camera : genesis.Camera The created camera object. """ + if denoise is None: + denoise = sys.platform != "darwin" return self._visualizer.add_camera( res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, env_idx ) diff --git a/genesis/ext/pyrender/offscreen.py b/genesis/ext/pyrender/offscreen.py index 0076d33738..c96188a239 100644 --- a/genesis/ext/pyrender/offscreen.py +++ b/genesis/ext/pyrender/offscreen.py @@ -173,16 +173,22 @@ def render( else: 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) + glBindFramebuffer(GL_READ_FRAMEBUFFER, 0) + glReadBuffer(GL_FRONT) if depth: - depth = renderer.read_depth_buf() + z_near = scene.main_camera_node.camera.znear + z_far = scene.main_camera_node.camera.zfar + if z_far is None: + z_far = -1.0 + depth_arr = renderer.jit.read_depth_buf(self.viewport_height, self.viewport_width, z_near, z_far) + depth_arr = renderer._resize_image(depth_arr, antialias=not seg) if flags & RenderFlags.DEPTH_ONLY: - retval = (depth,) + retval = (depth_arr,) else: - color = renderer.read_color_buf() - retval = (color, depth) + color_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False) + color_arr = renderer._resize_image(color_arr, antialias=not seg) + retval = (color_arr, depth_arr) if depth else (color_arr,) else: retval = () @@ -207,7 +213,16 @@ 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, force_skip_shadows=True) + + if self._platform.supports_framebuffers(): + normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True) + else: + glBindFramebuffer(GL_READ_FRAMEBUFFER, 0) + glReadBuffer(GL_FRONT) + renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True) + normal_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False) + normal_arr = renderer._resize_image(normal_arr, antialias=not seg) + retval = (*retval, normal_arr) renderer._program_cache = old_cache diff --git a/genesis/ext/pyrender/viewer.py b/genesis/ext/pyrender/viewer.py index 54dd11ab63..d21714ff3c 100644 --- a/genesis/ext/pyrender/viewer.py +++ b/genesis/ext/pyrender/viewer.py @@ -1063,11 +1063,12 @@ def _save_image(self): filename = self._get_save_filename(["png", "jpg", "gif", "all"]) if filename is not None: self.viewer_flags["save_directory"] = os.path.dirname(filename) - imageio.imwrite(filename, self._renderer.read_color_buf()) + data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False) + imageio.imwrite(filename, img_arr) def _record(self): """Save another frame for the GIF.""" - data = self._renderer.read_color_buf() + data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False) if not np.all(data == 0.0): self.video_recorder.write_frame(data) diff --git a/genesis/vis/camera.py b/genesis/vis/camera.py index b1924f6803..a93e2e2c8b 100644 --- a/genesis/vis/camera.py +++ b/genesis/vis/camera.py @@ -130,6 +130,13 @@ def __init__( self._focus_dist = np.linalg.norm(np.asarray(lookat) - np.asarray(pos)) def build(self): + n_envs = max(self._visualizer.scene.n_envs, 1) + self._multi_env_pos_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device) + self._multi_env_lookat_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device) + self._multi_env_up_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device) + self._multi_env_transform_tensor = torch.empty((n_envs, 4, 4), dtype=gs.tc_float, device=gs.device) + self._multi_env_quat_tensor = torch.empty((n_envs, 4), dtype=gs.tc_float, device=gs.device) + self._envs_offset = torch.as_tensor(self._visualizer._scene.envs_offset, dtype=gs.tc_float, device=gs.device) self._batch_renderer = self._visualizer.batch_renderer @@ -154,7 +161,12 @@ def build(self): self._other_stacked = self._visualizer._context.env_separate_rigid self._is_built = True - self.setup_initial_env_poses() + self.set_pose( + transform=self._initial_transform, pos=self._initial_pos, lookat=self._initial_lookat, up=self._initial_up + ) + # FIXME: For some reason, it is necessary to update the camera twice... + if self._raytracer is not None: + self._raytracer.update_camera(self) def attach(self, rigid_link, offset_T): """ @@ -296,8 +308,8 @@ def _batch_render( depth=False, segmentation=False, normal=False, + antialiasing=True, force_render=False, - antialiasing=False, ): """ Render the camera view with batch renderer. @@ -371,7 +383,7 @@ def render( normal_arr : np.ndarray The rendered surface normal(s). """ - rgb_arr, depth_arr, seg_arr, seg_color_arr, normal_arr = None, None, None, None, None + rgb_arr, depth_arr, seg_arr, seg_color_arr, seg_idxc_arr, normal_arr = None, None, None, None, None, None if self._batch_renderer is not None: rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._batch_render( @@ -521,28 +533,6 @@ def render_pointcloud(self, world_frame=True): point_cloud = point_cloud[:, :3].reshape((*depth_arr.shape, 3)) return point_cloud, mask - @gs.assert_built - def setup_initial_env_poses(self): - """ - Setup the camera poses for multiple environments. - """ - if self._initial_transform is not None: - assert self._initial_transform.shape == (4, 4) - self._initial_pos, self._initial_lookat, self._initial_up = gu.T_to_pos_lookat_up(self._initial_transform) - else: - self._initial_transform = gu.pos_lookat_up_to_T(self._initial_pos, self._initial_lookat, self._initial_up) - - n_envs = max(self._visualizer.scene.n_envs, 1) - self._multi_env_pos_tensor = self._initial_pos.repeat((n_envs, 1)) - self._multi_env_lookat_tensor = self._initial_lookat.repeat((n_envs, 1)) - self._multi_env_up_tensor = self._initial_up.repeat((n_envs, 1)) - self._multi_env_transform_tensor = self._initial_transform.repeat((n_envs, 1, 1)) - self._multi_env_quat_tensor = _T_to_quat_for_madrona(self._multi_env_transform_tensor) - - self._rasterizer.update_camera(self) - if self._raytracer is not None: - self._raytracer.update_camera(self) - @gs.assert_built def set_pose(self, transform=None, pos=None, lookat=None, up=None, env_idx=None): """ diff --git a/genesis/vis/rasterizer.py b/genesis/vis/rasterizer.py index 32c3386d57..724ef9ec35 100644 --- a/genesis/vis/rasterizer.py +++ b/genesis/vis/rasterizer.py @@ -69,36 +69,37 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma self._context.buffer.clear() # Render - if rgb or depth or normal: - retval = self._renderer.render( - self._context._scene, - self._camera_targets[camera.uid], - camera_node=self._camera_nodes[camera.uid], - env_separate_rigid=self._context.env_separate_rigid, - rgb=rgb, - normal=normal, - seg=False, - depth=depth, - plane_reflection=rgb and self._context.plane_reflection, - shadow=rgb and self._context.shadow, - ) - - if segmentation: - seg_idxc_rgb_arr, *_ = self._renderer.render( - self._context._scene, - self._camera_targets[camera.uid], - camera_node=self._camera_nodes[camera.uid], - env_separate_rigid=self._context.env_separate_rigid, - rgb=False, - normal=False, - seg=True, - depth=False, - plane_reflection=False, - shadow=False, - ) - - # Unset the context - self._renderer.make_uncurrent() + try: + if rgb or depth or normal: + retval = self._renderer.render( + self._context._scene, + self._camera_targets[camera.uid], + camera_node=self._camera_nodes[camera.uid], + env_separate_rigid=self._context.env_separate_rigid, + rgb=rgb, + normal=normal, + seg=False, + depth=depth, + plane_reflection=rgb and self._context.plane_reflection, + shadow=rgb and self._context.shadow, + ) + + if segmentation: + seg_idxc_rgb_arr, *_ = self._renderer.render( + self._context._scene, + self._camera_targets[camera.uid], + camera_node=self._camera_nodes[camera.uid], + env_separate_rigid=self._context.env_separate_rigid, + rgb=False, + normal=False, + seg=True, + depth=False, + plane_reflection=False, + shadow=False, + ) + finally: + # Unset the context + self._renderer.make_uncurrent() else: # Render if rgb or depth or normal: diff --git a/genesis/vis/raytracer.py b/genesis/vis/raytracer.py index 694b14712f..6ef6a89def 100644 --- a/genesis/vis/raytracer.py +++ b/genesis/vis/raytracer.py @@ -604,7 +604,7 @@ def add_camera(self, camera): if camera_model == "pinhole": self._cameras[camera_name] = LuisaRenderPy.PinholeCamera( name=camera_name, - pose=self.get_transform(camera.transform), + pose=self.get_transform(np.eye(4)), film=LuisaRenderPy.Film(resolution=camera.res), filter=LuisaRenderPy.Filter(), spp=camera.spp, @@ -613,7 +613,7 @@ def add_camera(self, camera): elif camera_model == "thinlens": self._cameras[camera_name] = LuisaRenderPy.ThinLensCamera( name=camera_name, - pose=self.get_transform(camera.transform), + pose=self.get_transform(np.eye(4)), film=LuisaRenderPy.Film(resolution=camera.res), filter=LuisaRenderPy.Filter(), spp=camera.spp, diff --git a/tests/test_render.py b/tests/test_render.py index 1aca0f16d2..c93ba04996 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -119,7 +119,7 @@ def test_segmentation(segmentation_level, particle_mode): @pytest.mark.required @pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin")) -def test_batched_offscreen_rendering(show_viewer, tol): +def test_batched_offscreen_rendering(tmp_path, show_viewer, tol): scene = gs.Scene( vis_options=gs.options.VisOptions( # rendered_envs_idx=(0, 1, 2), @@ -247,6 +247,7 @@ def test_batched_offscreen_rendering(show_viewer, tol): ) scene.build(n_envs=3, env_spacing=(2.0, 2.0)) + cam.start_recording() for _ in range(7): dofs_lower_bound, dofs_upper_bound = robot.get_dofs_limit() qpos = dofs_lower_bound + (dofs_upper_bound - dofs_lower_bound) * torch.rand(robot.n_qs) @@ -270,6 +271,7 @@ def test_batched_offscreen_rendering(show_viewer, tol): for i in range(3): assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol) + cam.stop_recording(save_to_filename=(tmp_path / "video.mp4")) @pytest.mark.required