diff --git a/examples/rigid/single_franka_batch_render.py b/examples/rigid/single_franka_batch_render.py index 285e2e12e0..1dc237d995 100644 --- a/examples/rigid/single_franka_batch_render.py +++ b/examples/rigid/single_franka_batch_render.py @@ -13,9 +13,10 @@ def main(): parser.add_argument("-b", "--n_envs", type=int, default=3) parser.add_argument("-s", "--n_steps", type=int, default=2) parser.add_argument("-r", "--render_all_cameras", action="store_true", default=False) - parser.add_argument("-o", "--output_dir", type=str, default="img_output/test") + parser.add_argument("-o", "--output_dir", type=str, default="data/test") parser.add_argument("-u", "--use_rasterizer", action="store_true", default=False) parser.add_argument("-d", "--debug", action="store_true", default=False) + parser.add_argument("-l", "--seg_level", type=str, default="link") args = parser.parse_args() ########################## init ########################## @@ -23,6 +24,9 @@ def main(): ########################## create a scene ########################## scene = gs.Scene( + vis_options=gs.options.VisOptions( + segmentation_level=args.seg_level, + ), renderer=gs.options.renderers.BatchRenderer( use_rasterizer=args.use_rasterizer, ), @@ -61,21 +65,31 @@ def main(): fov=45, GUI=args.vis, ) + cam_2 = scene.add_camera( + res=(512, 512), + pos=(0.0, 0.1, 5.0), + lookat=(0.0, 0.0, 0.0), + fov=45, + GUI=args.vis, + ) + scene.add_light( - pos=[0.0, 0.0, 1.5], - dir=[1.0, 1.0, -2.0], - directional=1, - castshadow=1, + pos=(0.0, 0.0, 1.5), + dir=(1.0, 1.0, -2.0), + color=(1.0, 0.0, 0.0), + directional=True, + castshadow=True, cutoff=45.0, intensity=0.5, ) scene.add_light( - pos=[4, -4, 4], - dir=[-1, 1, -1], - directional=0, - castshadow=1, - cutoff=45.0, - intensity=0.5, + pos=(4, -4, 4), + dir=(0, 0, -1), + directional=False, + castshadow=True, + cutoff=80.0, + intensity=1.0, + attenuation=0.1, ) ########################## build ########################## @@ -91,11 +105,19 @@ def main(): if args.debug: debug_cam.render() if args.render_all_cameras: - rgba, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True) - exporter.export_frame_all_cameras(i, rgb=rgba, depth=depth) + color, depth, seg, normal = scene.render_all_cameras( + rgb=True, depth=i % 2 == 1, segmentation=i % 2 == 1, normal=True + ) + exporter.export_frame_all_cameras(i, rgb=color, depth=depth, segmentation=seg, normal=normal) else: - rgba, depth, _, _ = cam_1.render(rgb=True, depth=True) - exporter.export_frame_single_camera(i, cam_1.idx, rgb=rgba, depth=depth) + color, depth, seg, normal = cam_1.render( + rgb=False, + depth=True, + segmentation=True, + colorize_seg=True, + normal=False, + ) + exporter.export_frame_single_camera(i, cam_1.idx, rgb=seg, depth=depth, segmentation=None, normal=normal) if args.debug: debug_cam.stop_recording("debug_cam.mp4") diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 80dbe2e2f0..521ad2ba89 100644 --- a/genesis/engine/scene.py +++ b/genesis/engine/scene.py @@ -451,77 +451,94 @@ def link_entities( parent_link._child_idxs.append(child_link.idx) @gs.assert_unbuilt - def add_light( + def add_mesh_light( self, - *, morph: Morph | None = None, color: ArrayLike | None = (1.0, 1.0, 1.0, 1.0), intensity: float = 20.0, revert_dir: bool | None = False, double_sided: bool | None = False, - beam_angle: float | None = 180.0, - pos: ArrayLike | None = None, - dir: ArrayLike | None = None, - directional: bool | None = None, - castshadow: bool | None = None, - cutoff: float | None = None, + cutoff: float | None = 180.0, ): """ - Add a light to the scene. - - Warning - ------- - The signature of this method is different depending on the renderer being used, i.e.: - - RayTracer: 'add_light(self, morph, color, intensity, revert_dir, double_sided, beam_angle)' - - BatchRender: 'add_ligth(self, pos, dir, intensity, directional, castshadow, cutoff)' - - Rasterizer: **Unsupported** + Add a mesh light to the scene. Only supported by RayTracer. Parameters ---------- morph : gs.morphs.Morph - The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`. Only supported by - RayTracer. + The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`. color : tuple of float, shape (3,) - The color of the light, specified as (r, g, b). Only supported by RayTracer. + The color of the light, specified as (r, g, b). intensity : float The intensity of the light. revert_dir : bool Whether to revert the direction of the light. If True, the light will be emitted towards the mesh's inside. - Only supported by RayTracer. double_sided : bool - Whether to emit light from both sides of surface. Only supported by RayTracer. - beam_angle : float - The beam angle of the light. Only supported by RayTracer. + Whether to emit light from both sides of surface. + cutoff : float + The cutoff angle of the light in degrees. Range: [0.0, 180.0]. + """ + if not isinstance(self.renderer_options, gs.renderers.RayTracer): + if isinstance(self.renderer_options, gs.renderers.BatchRenderer): + gs.raise_exception( + "This method is only supported by RayTracer. Please use 'add_light' when using BatchRenderer." + ) + else: + gs.raise_exception( + "This method is only supported by RayTracer. Impossible to add light when using Rasterizer." + ) + + if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)): + gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.") + mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False)) + self._visualizer.add_mesh_light(mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, cutoff) + + @gs.assert_unbuilt + def add_light( + self, + pos: ArrayLike, + dir: ArrayLike, + color: ArrayLike = (1.0, 1.0, 1.0), + intensity: float = 1.0, + directional: bool = False, + castshadow: bool = True, + cutoff: float = 45.0, + attenuation: float = 0.0, + ): + """ + Add a light to the scene for batch renderer. + + Parameters + ---------- pos : tuple of float, shape (3,) - The position of the light, specified as (x, y, z). Only supported by BatchRenderer. + The position of the light, specified as (x, y, z). dir : tuple of float, shape (3,) - The direction of the light, specified as (x, y, z). Only supported by BatchRenderer. + The direction of the light, specified as (x, y, z). + color : tuple of float, shape (3,) + The color of the light, specified as (r, g, b). intensity : float - The intensity of the light. Only supported by BatchRenderer. + The intensity of the light. directional : bool - Whether the light is directional. Only supported by BatchRenderer. + Whether the light is directional. castshadow : bool - Whether the light casts shadows. Only supported by BatchRenderer. + Whether the light casts shadows. cutoff : float - The cutoff angle of the light in degrees. Only supported by BatchRenderer. + The cutoff angle of the light in degrees. Range: (0.0, 90.0). + attenuation : float + The attenuation factor of the light. + Light intensity will attenuate by distance with (1 / (1 + attenuation * distance ^ 2)) """ - if self._visualizer.batch_renderer is not None: - if any(map(lambda e: e is None, (pos, dir, intensity, directional, castshadow, cutoff))): - gs.raise_exception("Input arguments do not complain with expected signature when using 'BatchRenderer'") - - self.visualizer.add_light(pos, dir, intensity, directional, castshadow, cutoff) - elif self.visualizer.raytracer is not None: - if any(map(lambda e: e is None, (morph, color, intensity, revert_dir, double_sided, beam_angle))): - gs.raise_exception("Input arguments do not complain with expected signature when using 'RayTracer'") - if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)): - gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.") - - mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False)) - self.visualizer.raytracer.add_mesh_light( - mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, beam_angle - ) - else: - gs.raise_exception("Adding lights is only supported by 'RayTracer' and 'BatchRenderer'.") + if not isinstance(self.renderer_options, gs.renderers.BatchRenderer): + if isinstance(self.renderer_options, gs.renderers.BatchRenderer): + gs.raise_exception( + "This method is only supported by BatchRenderer. Please use 'add_mesh_light' when using RayTracer." + ) + else: + gs.raise_exception( + "This method is only supported by BatchRenderer. Impossible to add light when using Rasterizer." + ) + + self.visualizer.add_light(pos, dir, color, intensity, directional, castshadow, cutoff, attenuation) @gs.assert_unbuilt def add_sensor(self, sensor_options: "SensorOptions"): @@ -541,6 +558,8 @@ def add_camera( GUI=False, spp=256, denoise=None, + near=0.05, + far=100.0, env_idx=None, debug=False, ): @@ -581,6 +600,12 @@ def add_camera( Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults 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. + near : float + Distance from camera center to near plane in meters. + Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 0.05. + far : float + Distance from camera center to far plane in meters. + Only available when using rasterizer in Rasterizer and BatchRender renderer. Defaults to 100.0. env_idx : int, optional The specific environment index to bind to the camera. This option must be specified if and only if a non-batched renderer is being used. If provided, only this environment will be taken into account when @@ -602,7 +627,7 @@ def add_camera( 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, debug + res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, near, far, env_idx, debug ) @gs.assert_unbuilt @@ -1120,7 +1145,16 @@ def draw_debug_path(self, qposs, entity, link_idx=-1, density=0.3, frame_scaling ) @gs.assert_built - def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=False, force_render=False): + def render_all_cameras( + self, + rgb=True, + depth=False, + segmentation=False, + colorize_seg=False, + normal=False, + antialiasing=False, + force_render=False, + ): """ Render the scene for all cameras using the batch renderer. @@ -1130,10 +1164,12 @@ def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=F Whether to render the rgb image. depth : bool, optional Whether to render the depth image. - normal : bool, optional - Whether to render the normal image. segmentation : bool, optional Whether to render the segmentation image. + normal : bool, optional + Whether to render the normal image. + antialiasing : bool, optional + Whether to apply anti-aliasing. force_render : bool, optional Whether to force render the scene. @@ -1145,7 +1181,12 @@ def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=F if self._visualizer.batch_renderer is None: gs.raise_exception("Method only supported by 'BatchRenderer'") - return self._visualizer.batch_renderer.render(rgb, depth, normal, segmentation, force_render) + rgb_out, depth_out, seg_out, normal_out = self._visualizer.batch_renderer.render( + rgb, depth, segmentation, normal, antialiasing, force_render + ) + if segmentation and colorize_seg: + seg_out = tuple(self._visualizer.batch_renderer.colorize_seg_idxc_arr(seg) for seg in seg_out) + return rgb_out, depth_out, seg_out, normal_out @gs.assert_built def clear_debug_object(self, object): @@ -1396,3 +1437,18 @@ def fem_solver(self): def pbd_solver(self): """The scene's `pbd_solver`, managing all the `PBDEntity` in the scene.""" return self._sim.pbd_solver + + @property + def segmentation_idx_dict(self): + """ + Returns a dictionary mapping segmentation indices to scene entities. + + In the segmentation map: + - Index 0 corresponds to the background (-1). + - Indices > 0 correspond to scene elements, which may be represented as: + - `entity_id` + - `(entity_id, link_id)` + - `(entity_id, link_id, geom_id)` + depending on the material type and the configured segmentation level. + """ + return self._visualizer.segmentation_idx_dict diff --git a/genesis/options/surfaces.py b/genesis/options/surfaces.py index 9f857e26dc..144db58d1c 100644 --- a/genesis/options/surfaces.py +++ b/genesis/options/surfaces.py @@ -67,8 +67,8 @@ class Surface(Options): double_sided : bool | None, optional Whether to render both sides of the surface. Useful for non-watertight 2D objects. Defaults to True for Cloth material and False for others. - beam_angle : float - The beam angle of emission. Defaults to 180.0. + cutoff : float + The cutoff angle of emission. Defaults to 180.0. normal_diff_clamp : float, optional Controls the threshold for computing surface normals by interpolating vertex normals. recon_backend : str, optional @@ -99,8 +99,8 @@ class Surface(Options): vis_mode: Optional[str] = None smooth: bool = True double_sided: Optional[bool] = None - beam_angle: float = 180 - normal_diff_clamp: float = 180 + cutoff: float = 180.0 + normal_diff_clamp: float = 180.0 recon_backend: str = "splashsurf" generate_foam: bool = False foam_options: Optional[FoamOptions] = None diff --git a/genesis/utils/image_exporter.py b/genesis/utils/image_exporter.py index 0339d1c397..a04e846592 100644 --- a/genesis/utils/image_exporter.py +++ b/genesis/utils/image_exporter.py @@ -194,15 +194,12 @@ def export_frame_single_camera( imgs_data = np.asarray(imgs_data) # Make sure that image data has shape `(n_env, H, W [, C>1])`` - is_single_channel = img_type in (IMAGE_TYPE.DEPTH, IMAGE_TYPE.SEGMENTATION) - if imgs_data.ndim == (2 if is_single_channel else 3): + if imgs_data.ndim < 4: imgs_data = imgs_data[None] - if imgs_data.ndim == 4 and is_single_channel: + if imgs_data.ndim == 4 and imgs_data.shape[-1] == 1: imgs_data = imgs_data[..., 0] - if is_single_channel and imgs_data.ndim != 3: - gs.raise_exception("'{imgs_data}' images must be tensors of shape (n_envs, H, W)") - elif not is_single_channel and (imgs_data.ndim != 4 or imgs_data.shape[-1] != 3): - gs.raise_exception("'{imgs_data}' images must be tensors of shape (n_envs, H, W, 3)") + if imgs_data.ndim not in (3, 4): + gs.raise_exception("'{imgs_data}' images must be tensors of shape (n_envs, H, W [, C>1])") # Convert image data to grayscale array if necessary if img_type == IMAGE_TYPE.DEPTH: @@ -214,7 +211,7 @@ def export_frame_single_camera( imgs_data = imgs_data.astype(np.uint8) # Flip channel order if necessary - if not is_single_channel: + if imgs_data.ndim == 4: imgs_data = np.flip(imgs_data, axis=-1) # Export image array as (compressed) PNG file. diff --git a/genesis/vis/batch_renderer.py b/genesis/vis/batch_renderer.py index 201478a7c7..f31a20fbbe 100644 --- a/genesis/vis/batch_renderer.py +++ b/genesis/vis/batch_renderer.py @@ -2,13 +2,18 @@ import numpy as np import torch +from PIL import Image +from trimesh.visual.texture import TextureVisuals +from trimesh.visual.color import ColorVisuals import genesis as gs from genesis.repr_base import RBC from genesis.constants import IMAGE_TYPE +from .rasterizer_context import SegmentationColorMap + try: - from gs_madrona.renderer_gs import MadronaBatchRendererAdapter + from gs_madrona.renderer_gs import MadronaBatchRendererAdapter, GeomRetriever except ImportError as e: gs.raise_exception_from("Madrona batch renderer is only supported on Linux x86-64.", e) @@ -23,14 +28,178 @@ def _make_tensor(data, *, dtype: torch.dtype = torch.float32): return torch.tensor(data, dtype=dtype, device=gs.device) +class GenesisGeomRetriever(GeomRetriever): + def __init__(self, rigid_solver, seg_level): + self.rigid_solver = rigid_solver + self.seg_color_map = SegmentationColorMap(to_torch=True) + self.seg_level = seg_level + self.geom_idxc = None + + self.default_geom_group = 2 + self.default_enabled_geom_groups = np.array([self.default_geom_group], dtype=np.int32) + self.num_textures_per_material = 10 # Madrona allows up to 10 textures per material, should not change. + + def build(self): + self.n_vgeoms = self.rigid_solver.n_vgeoms + self.geom_idxc = [] + vgeoms = self.rigid_solver.vgeoms + for vgeom in vgeoms: + seg_key = self.get_seg_key(vgeom) + seg_idxc = self.seg_color_map.seg_key_to_idxc(seg_key) + self.geom_idxc.append(seg_idxc) + self.geom_idxc = torch.tensor(self.geom_idxc, dtype=torch.int32, device=gs.device) + self.seg_color_map.generate_seg_colors() + + def get_seg_key(self, vgeom): + if self.seg_level == "geom": + return (vgeom.entity.idx, vgeom.link.idx, vgeom.idx) + elif self.seg_level == "link": + return (vgeom.entity.idx, vgeom.link.idx) + elif self.seg_level == "entity": + return vgeom.entity.idx + else: + gs.raise_exception(f"Unsupported segmentation level: {self.seg_level}") + + # FIXME: Use a kernel to do it efficiently + def retrieve_rigid_meshes_static(self): + args = {} + vgeoms = self.rigid_solver.vgeoms + + # Retrieve geom data + mesh_vertices = self.rigid_solver.vverts_info.init_pos.to_numpy() + mesh_faces = self.rigid_solver.vfaces_info.vverts_idx.to_numpy() + mesh_vertex_offsets = self.rigid_solver.vgeoms_info.vvert_start.to_numpy() + mesh_face_starts = self.rigid_solver.vgeoms_info.vface_start.to_numpy() + mesh_face_ends = self.rigid_solver.vgeoms_info.vface_end.to_numpy() + for i in range(self.n_vgeoms): + mesh_faces[mesh_face_starts[i] : mesh_face_ends[i]] -= mesh_vertex_offsets[i] + + geom_data_ids = [] + for vgeom in vgeoms: + seg_key = self.get_seg_key(vgeom) + seg_id = self.seg_color_map.seg_key_to_idxc(seg_key) + geom_data_ids.append(seg_id) + + args["mesh_vertices"] = mesh_vertices + args["mesh_faces"] = mesh_faces + args["mesh_vertex_offsets"] = mesh_vertex_offsets + args["mesh_face_offsets"] = mesh_face_starts + args["geom_types"] = np.full((self.n_vgeoms,), 7, dtype=np.int32) # 7 stands for mesh + args["geom_groups"] = np.full((self.n_vgeoms,), self.default_geom_group, dtype=np.int32) + args["geom_data_ids"] = np.arange(self.n_vgeoms, dtype=np.int32) + args["geom_sizes"] = np.ones((self.n_vgeoms, 3), dtype=np.float32) + args["enabled_geom_groups"] = self.default_enabled_geom_groups + + # Retrieve material data + num_materials = 0 + total_uv_size = 0 + total_texture_size = 0 + geom_mat_ids = [] + geom_uv_sizes = [] + geom_uv_offsets = [] + geom_rgbas = [] + + mat_uv_data = [] + mat_texture_widths = [] + mat_texture_heights = [] + mat_texture_nchans = [] + mat_texture_offsets = [] + mat_texture_data = [] + mat_texture_ids = [] + mat_rgbas = [] + + for vgeom in vgeoms: + visual = vgeom.get_trimesh().visual + if isinstance(visual, TextureVisuals): + uv_size = visual.uv.shape[0] + geom_mat_ids.append(num_materials) + geom_uv_sizes.append(uv_size) + geom_uv_offsets.append(total_uv_size) + geom_rgbas.append(np.zeros((4,), dtype=np.float32)) + + texture_width = visual.material.image.width + texture_height = visual.material.image.height + texture_nchans = 4 if visual.material.image.mode == "RGBA" else 3 + texture_size = texture_width * texture_height * texture_nchans + texture_ids = np.full((self.num_textures_per_material,), -1, np.int32) + texture_ids[0] = num_materials + + mat_uv_data.append(visual.uv.astype(np.float32)) + mat_texture_widths.append(texture_width) + mat_texture_heights.append(texture_height) + mat_texture_nchans.append(texture_nchans) + mat_texture_offsets.append(total_texture_size) + mat_texture_data.append( + np.asarray(visual.material.image.transpose(Image.Transpose.FLIP_TOP_BOTTOM), dtype=np.uint8).flat + ) + mat_texture_ids.append(texture_ids) + mat_rgbas.append(visual.material.diffuse.astype(np.float32) / 255.0) + + num_materials += 1 + total_uv_size += uv_size + total_texture_size += texture_size + else: + geom_mat_ids.append(-1) + geom_uv_sizes.append(0) + geom_uv_offsets.append(-1) + if isinstance(visual, ColorVisuals): + geom_rgbas.append(visual.main_color.astype(np.float32) / 255.0) + else: + geom_rgbas.append(np.zeros((4,), dtype=np.float32)) + + args["geom_mat_ids"] = np.array(geom_mat_ids, np.int32) + args["mesh_texcoord_num"] = np.array(geom_uv_sizes, np.int32) + args["mesh_texcoord_offsets"] = np.array(geom_uv_offsets, np.int32) + args["geom_rgba"] = np.stack(geom_rgbas, axis=0) + + args["mesh_texcoords"] = np.concatenate(mat_uv_data, axis=0) if mat_uv_data else np.empty((0, 2), np.float32) + args["tex_widths"] = np.array(mat_texture_widths, np.int32) + args["tex_heights"] = np.array(mat_texture_heights, np.int32) + args["tex_nchans"] = np.array(mat_texture_nchans, np.int32) + args["tex_offsets"] = np.array(mat_texture_offsets, np.int32) + args["tex_data"] = np.concatenate(mat_texture_data, axis=0) if mat_texture_data else np.array([], np.uint8) + args["mat_tex_ids"] = ( + np.stack(mat_texture_ids, axis=0) + if mat_texture_ids + else np.empty((0, self.num_textures_per_material), np.int32) + ) + args["mat_rgba"] = np.stack(mat_rgbas, axis=0) if mat_rgbas else np.empty((0, 4), np.float32) + + return args + + # FIXME: Use a kernel to do it efficiently + def retrieve_rigid_property_torch(self, num_worlds): + geom_rgb_torch = self.rigid_solver.vgeoms_info.color.to_torch() + geom_rgb_int = (geom_rgb_torch * 255).to(torch.int32) + geom_rgb_uint = (geom_rgb_int[:, 0] << 16) | (geom_rgb_int[:, 1] << 8) | geom_rgb_int[:, 2] + geom_rgb = geom_rgb_uint.unsqueeze(0).repeat(num_worlds, 1) + + geom_mat_ids = torch.full((self.n_vgeoms,), -1, dtype=torch.int32, device=gs.device) + geom_mat_ids = geom_mat_ids.unsqueeze(0).repeat(num_worlds, 1) + + geom_sizes = torch.ones((self.n_vgeoms, 3), dtype=torch.float32, device=gs.device) + geom_sizes = geom_sizes.unsqueeze(0).repeat(num_worlds, 1, 1) + return geom_mat_ids, geom_rgb, geom_sizes + + # FIXME: Use a kernel to do it efficiently + def retrieve_rigid_state_torch(self): + geom_pos = self.rigid_solver.vgeoms_state.pos.to_torch() + geom_rot = self.rigid_solver.vgeoms_state.quat.to_torch() + geom_pos = geom_pos.transpose(0, 1).contiguous() + geom_rot = geom_rot.transpose(0, 1).contiguous() + return geom_pos, geom_rot + + class Light: - def __init__(self, pos, dir, intensity, directional, castshadow, cutoff): + def __init__(self, pos, dir, color, intensity, directional, castshadow, cutoff, attenuation): self._pos = pos self._dir = tuple(dir / np.linalg.norm(dir)) + self._color = color self._intensity = intensity self._directional = directional self._castshadow = castshadow self._cutoff = cutoff + self._attenuation = attenuation @property def pos(self): @@ -40,6 +209,10 @@ def pos(self): def dir(self): return self._dir + @property + def color(self): + return self._color + @property def intensity(self): return self._intensity @@ -60,22 +233,27 @@ def cutoffRad(self): def cutoffDeg(self): return self._cutoff + @property + def attenuation(self): + return self._attenuation + class BatchRenderer(RBC): """ This class is used to manage batch rendering """ - def __init__(self, visualizer, renderer_options): + def __init__(self, visualizer, renderer_options, vis_options): self._visualizer = visualizer self._lights = gs.List() self._use_rasterizer = renderer_options.use_rasterizer self._renderer = None + self._geom_retriever = GenesisGeomRetriever(self._visualizer.scene.rigid_solver, vis_options.segmentation_level) self._data_cache = {} self._t = -1 - def add_light(self, pos, dir, intensity, directional, castshadow, cutoff): - self._lights.append(Light(pos, dir, intensity, directional, castshadow, cutoff)) + def add_light(self, pos, dir, color, intensity, directional, castshadow, cutoff, attenuation): + self._lights.append(Light(pos, dir, color, intensity, directional, castshadow, cutoff, attenuation)) def build(self): """ @@ -90,6 +268,9 @@ def build(self): if not self._cameras: gs.raise_exception("Please add at least one camera when using BatchRender.") + # Build the geometry retriever + self._geom_retriever.build() + # Make sure that all cameras have identical resolution try: ((camera_width, camera_height),) = set(camera.res for camera in self._cameras) @@ -97,33 +278,36 @@ def build(self): gs.raise_exception_from("All cameras must have the exact same resolution when using BatchRender.", e) self._renderer = MadronaBatchRendererAdapter( - rigid=self._visualizer.scene.rigid_solver, + geom_retriever=self._geom_retriever, gpu_id=gs.device.index if gs.device.index is not None else 0, num_worlds=max(self._visualizer.scene.n_envs, 1), num_cameras=len(self._cameras), num_lights=len(self._lights), cam_fovs_tensor=_make_tensor([camera.fov for camera in self._cameras]), + cam_znears_tensor=_make_tensor([camera.near for camera in self._cameras]), + cam_zfars_tensor=_make_tensor([camera.far for camera in self.cameras]), batch_render_view_width=camera_width, batch_render_view_height=camera_height, add_cam_debug_geo=False, use_rasterizer=self._use_rasterizer, ) self._renderer.init( - rigid=self._visualizer.scene.rigid_solver, cam_pos_tensor=torch.stack([camera.get_pos() for camera in self._cameras], dim=1), cam_rot_tensor=_transform_camera_quat(torch.stack([camera.get_quat() for camera in self._cameras], dim=1)), lights_pos_tensor=_make_tensor([light.pos for light in self._lights]).reshape((-1, 3)), lights_dir_tensor=_make_tensor([light.dir for light in self._lights]).reshape((-1, 3)), - lights_intensity_tensor=_make_tensor([light.intensity for light in self._lights]), + lights_rgb_tensor=_make_tensor([light.color for light in self._lights]).reshape((-1, 3)), lights_directional_tensor=_make_tensor([light.directional for light in self._lights], dtype=torch.bool), lights_castshadow_tensor=_make_tensor([light.castshadow for light in self._lights], dtype=torch.bool), lights_cutoff_tensor=_make_tensor([light.cutoffRad for light in self._lights]), + lights_attenuation_tensor=_make_tensor([light.attenuation for light in self._lights]), + lights_intensity_tensor=_make_tensor([light.intensity for light in self._lights]), ) def update_scene(self): self._visualizer._context.update() - def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_render=False, aliasing=False): + def render(self, rgb=True, depth=False, segmentation=False, normal=False, antialiasing=False, force_render=False): """ Render all cameras in the batch. @@ -137,22 +321,22 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_ Whether to render the segmentation image. normal : bool, optional Whether to render the normal image. + antialiasing : bool, optional + Whether to apply anti-aliasing. force_render : bool, optional Whether to force render the scene. - aliasing : bool, optional - Whether to apply anti-aliasing. Returns ------- - rgb_arr : tuple of tensors + rgb_arr : tuple of arrays The sequence of rgb images associated with each camera. - depth_arr : tuple of tensors + depth_arr : tuple of arrays The sequence of depth images associated with each camera. + segmentation_arr : tuple of arrays + The sequence of segmentation images associated with each camera. + normal_arr : tuple of arrays + The sequence of normal images associated with each camera. """ - if normal: - raise NotImplementedError("Normal rendering not supported from now") - if segmentation: - raise NotImplementedError("Segmentation rendering not supported from now") # Clear cache if requested or necessary if force_render or self._t < self._visualizer.scene.t: @@ -160,7 +344,7 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_ # Fetch available cached data request = (rgb, depth, segmentation, normal) - cache_key = (aliasing,) + cache_key = (antialiasing,) cached = [self._data_cache.get((img_type, cache_key), None) for img_type in IMAGE_TYPE] # Force disabling rendering whenever cached data is already available @@ -183,13 +367,18 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_ needed[img_type] for img_type in (IMAGE_TYPE.RGB, IMAGE_TYPE.DEPTH, IMAGE_TYPE.NORMAL, IMAGE_TYPE.SEGMENTATION) ), - aliasing, + antialiasing, ), dtype=np.uint32, ) - rendered = list( - self._renderer.render(self._visualizer.scene.rigid_solver, cameras_pos, cameras_quat, render_flags) - ) + rendered = list(self._renderer.render(cameras_pos, cameras_quat, render_flags)) + + # convert seg geom idx to seg_idxc + if needed[IMAGE_TYPE.SEGMENTATION]: + seg_geoms = rendered[IMAGE_TYPE.SEGMENTATION] + mask = seg_geoms != -1 + seg_geoms[mask] = self._geom_retriever.geom_idxc[seg_geoms[mask]] + seg_geoms[~mask] = 0 # Post-processing: # * Remove alpha channel from RGBA @@ -202,6 +391,13 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_ data = data.squeeze(1) rendered[img_type] = tuple(data[..., :3].squeeze(-1)) + # Convert center distance depth to plane distance + if not self._use_rasterizer and needed[IMAGE_TYPE.DEPTH]: + rendered[IMAGE_TYPE.DEPTH] = tuple( + camera.distance_center_to_plane(depth_data) + for camera, depth_data in zip(self._cameras, rendered[IMAGE_TYPE.DEPTH]) + ) + # Update cache self._t = self._visualizer.scene.t for img_type, data in enumerate(rendered): @@ -211,6 +407,9 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_ # Return in the required order, or None if not requested return tuple(self._data_cache[(img_type, cache_key)] if needed[img_type] else None for img_type in IMAGE_TYPE) + def colorize_seg_idxc_arr(self, seg_idxc_arr): + return self._geom_retriever.seg_color_map.colorize_seg_idxc_arr(seg_idxc_arr) + def destroy(self): self._lights.clear() self._data_cache.clear() @@ -228,3 +427,7 @@ def lights(self): @property def cameras(self): return self._cameras + + @property + def seg_idxc_map(self): + return self._geom_retriever.seg_color_map.idxc_map diff --git a/genesis/vis/camera.py b/genesis/vis/camera.py index 3e9f01fadb..8e43400092 100644 --- a/genesis/vis/camera.py +++ b/genesis/vis/camera.py @@ -370,8 +370,8 @@ def render( segmentation=False, colorize_seg=False, normal=False, + antialiasing=False, force_render=False, - antialiasing=True, ): """ Render the camera view. @@ -400,10 +400,10 @@ def render( If True, the segmentation mask will be colorized. normal : bool, optional Whether to render the surface normal. - force_render : bool, optional - Whether to force rendering even if the scene has not changed. antialiasing : bool, optional Whether to apply anti-aliasing. Only supported by 'BatchRenderer' for now. + force_render : bool, optional + Whether to force rendering even if the scene has not changed. Returns ------- @@ -424,7 +424,7 @@ def render( 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( - rgb, depth, segmentation, normal, force_render, antialiasing + rgb_, depth, segmentation, normal, antialiasing, force_render ) elif self._raytracer is not None: if rgb_: @@ -480,20 +480,27 @@ def render( return rgb_arr if rgb else None, depth_arr, seg_arr, normal_arr - @gs.assert_built - def get_segmentation_idx_dict(self): - """ - Returns a dictionary mapping segmentation indices to scene entities. - - In the segmentation map: - - Index 0 corresponds to the background (-1). - - Indices > 0 correspond to scene elements, which may be represented as: - - `entity_id` - - `(entity_id, link_id)` - - `(entity_id, link_id, geom_id)` - depending on the material type and the configured segmentation level. - """ - return self._rasterizer._context.seg_idxc_map + def distance_center_to_plane(self, center_dis): + width, height = self.res + fx = fy = self.f + cx = self.cx + cy = self.cy + + if isinstance(center_dis, np.ndarray): + v, u = np.meshgrid(np.arange(height, dtype=np.int32), np.arange(width, dtype=np.int32), indexing="ij") + xd = (u + 0.5 - cx) / fx + yd = (v + 0.5 - cy) / fy + scale_inv = 1.0 / np.sqrt(xd**2 + yd**2 + 1.0) + else: # torch.Tensor + v, u = torch.meshgrid( + torch.arange(height, dtype=torch.int32, device=gs.device), + torch.arange(width, dtype=torch.int32, device=gs.device), + indexing="ij", + ) + xd = (u + 0.5 - cx) / fx + yd = (v + 0.5 - cy) / fy + scale_inv = torch.rsqrt(xd**2 + yd**2 + 1.0) + return center_dis * scale_inv @gs.assert_built def render_pointcloud(self, world_frame=True): @@ -513,12 +520,16 @@ def render_pointcloud(self, world_frame=True): mask_arr : np.ndarray The valid depth mask. boolean array of same shape as depth_arr """ - # Compute the (denormalized) depth map using PyRender systematically - # TODO: Add support of BatchRendered. - self._rasterizer.update_scene() - rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._rasterizer.render_camera( - self, rgb=False, depth=True, segmentation=False, normal=False - ) + # Compute the (denormalized) depth map + if self._batch_renderer is not None: + _, depth_arr, _, _ = self._batch_render(rgb=False, depth=True, segmentation=False, normal=False) + # FIXME: Avoid converting to numpy + depth_arr = tensor_to_array(depth_arr) + else: + self._rasterizer.update_scene() + _, depth_arr, _, _ = self._rasterizer.render_camera( + self, rgb=False, depth=True, segmentation=False, normal=False + ) # Convert OpenGL projection matrix to camera intrinsics width, height = self.res diff --git a/genesis/vis/rasterizer_context.py b/genesis/vis/rasterizer_context.py index 5b72288f53..e6c767cc5e 100644 --- a/genesis/vis/rasterizer_context.py +++ b/genesis/vis/rasterizer_context.py @@ -1,4 +1,5 @@ import numpy as np +import torch import trimesh import gstaichi as ti @@ -13,6 +14,69 @@ from genesis.utils.misc import tensor_to_array +class SegmentationColorMap: + def __init__(self, seed: int = 0, to_torch: bool = False): + self.seed = seed + self.to_torch = to_torch + self.idxc_map = {0: -1} + self.key_map = {-1: 0} + self.idxc_to_color = None + + def seg_key_to_idxc(self, seg_key): + seg_idxc = self.key_map.setdefault(seg_key, len(self.key_map)) + self.idxc_map[seg_idxc] = seg_key + return seg_idxc + + def seg_idxc_to_key(self, seg_idxc): + return self.idxc_map[seg_idxc] + + def colorize_seg_idxc_arr(self, seg_idxc_arr): + return self.idxc_to_color[seg_idxc_arr] + + def generate_seg_colors(self): + # seg_key: same as entity/link/geom's idx + # seg_idxc: segmentation index of objects + # seg_idxc_rgb: colorized seg_idxc internally used by renderer + + # Evenly spaced hues + num_keys = len(self.key_map) + hues = np.linspace(0.0, 1.0, num_keys, endpoint=False) + rng = np.random.default_rng(seed=self.seed) + rng.shuffle(hues) + + # Fixed saturation/value + s, v = 0.8, 0.95 + + # HSV to RGB conversion + rgb = np.zeros((num_keys, 3), dtype=np.float32) + i = (hues * 6).astype(np.int32) + f = hues * 6 - i + p = v * (1 - s) + q = v * (1 - f * s) + t = v * (1 - (1 - f) * s) + for k in range(1, num_keys): # Skip first color to enforce black background + match i[k] % 6: + case 0: + rgb[k] = (v, t[k], p) + case 1: + rgb[k] = (q[k], v, p) + case 2: + rgb[k] = (p, v, t[k]) + case 3: + rgb[k] = (p, q[k], v) + case 4: + rgb[k] = (t[k], p, v) + case 5: + rgb[k] = (v, p, q[k]) + rgb = np.round(rgb * 255.0).astype(np.uint8) + + # Store the generated map + if self.to_torch: + self.idxc_to_color = torch.from_numpy(rgb).to(device=gs.device) + else: + self.idxc_to_color = rgb + + class RasterizerContext: def __init__(self, options): self.show_world_frame = options.show_world_frame @@ -46,8 +110,7 @@ def __init__(self, options): self.dynamic_nodes = list() # nodes that live within single frame self.external_nodes = dict() # nodes added by external user self.seg_node_map = dict() - self.seg_idxc_map = {0: -1} - self.seg_key_map = {-1: 0} + self.seg_color_map = SegmentationColorMap() self.init_meshes() @@ -108,7 +171,7 @@ def build(self, scene): self.on_fem() # segmentation mapping - self.generate_seg_vars() + self.seg_color_map.generate_seg_colors() def destroy(self): self.clear_dynamic_nodes() @@ -140,7 +203,6 @@ def add_rigid_node(self, geom, obj, **kwargs): # create segemtation id if self.segmentation_level == "geom": seg_key = (geom.entity.idx, geom.link.idx, geom.idx) - assert False, "geom level segmentation not supported yet" elif self.segmentation_level == "link": seg_key = (geom.entity.idx, geom.link.idx) elif self.segmentation_level == "entity": @@ -870,27 +932,13 @@ def add_light(self, light): gs.raise_exception(f"Unsupported light type: {light['type']}") def create_node_seg(self, seg_key, seg_node): - seg_idxc = self.seg_key_to_idxc(seg_key) + seg_idxc = self.seg_color_map.seg_key_to_idxc(seg_key) if seg_node: self.seg_node_map[seg_node] = self.seg_idxc_to_idxc_rgb(seg_idxc) def remove_node_seg(self, seg_node): self.seg_node_map.pop(seg_node, None) - def generate_seg_vars(self): - # seg_key: same as entity/link/geom's idx - # seg_idxc: segmentation index of objects - # seg_idxc_rgb: colorized seg_idxc internally used by renderer - num_keys = len(self.seg_key_map) - rng = np.random.default_rng(seed=42) - self.seg_idxc_to_color = rng.integers(0, 255, size=(num_keys, 3), dtype=np.uint8) - self.seg_idxc_to_color[0] = 0 # background uses black - - def seg_key_to_idxc(self, seg_key): - seg_idxc = self.seg_key_map.setdefault(seg_key, len(self.seg_key_map)) - self.seg_idxc_map[seg_idxc] = seg_key - return seg_idxc - def seg_idxc_to_idxc_rgb(self, seg_idxc): seg_idxc_rgb = np.array( [ @@ -902,17 +950,18 @@ def seg_idxc_to_idxc_rgb(self, seg_idxc): ) return seg_idxc_rgb - def seg_idxc_to_key(self, seg_idxc): - return self.seg_idxc_map[seg_idxc] - def seg_idxc_rgb_arr_to_idxc_arr(self, seg_idxc_rgb_arr): # Combine the RGB components into a single integer seg_idxc_rgb_arr = seg_idxc_rgb_arr.astype(np.int64, copy=False) return seg_idxc_rgb_arr[..., 0] * (256 * 256) + seg_idxc_rgb_arr[..., 1] * 256 + seg_idxc_rgb_arr[..., 2] def colorize_seg_idxc_arr(self, seg_idxc_arr): - return self.seg_idxc_to_color[seg_idxc_arr] + return self.seg_color_map.colorize_seg_idxc_arr(seg_idxc_arr) @property def cameras(self): return self.visualizer.cameras + + @property + def seg_idxc_map(self): + return self.seg_color_map.idxc_map diff --git a/genesis/vis/raytracer.py b/genesis/vis/raytracer.py index 7e15f4e256..a1722b32f1 100644 --- a/genesis/vis/raytracer.py +++ b/genesis/vis/raytracer.py @@ -171,7 +171,7 @@ def __init__(self, options, vis_options): log_level=logging_class[self.logging_level], ) - def add_mesh_light(self, mesh, color, intensity, pos, quat, revert_dir=False, double_sided=False, beam_angle=180.0): + def add_mesh_light(self, mesh, color, intensity, pos, quat, revert_dir=False, double_sided=False, cutoff=180.0): color = np.array(color) if color.ndim != 1 or (color.shape[0] != 3 and color.shape[0] != 4): gs.raise_exception("Light color should have shape (3,) or (4,).") @@ -188,7 +188,7 @@ def add_mesh_light(self, mesh, color, intensity, pos, quat, revert_dir=False, do color[2] * intensity, ), double_sided=double_sided, - beam_angle=beam_angle, + cutoff=cutoff, ), name=str(mesh.uid), revert_dir=revert_dir, @@ -367,7 +367,7 @@ def add_surface(self, shape_name, surface): name=f"emis_{shape_name}", emission=self.get_texture(surface.get_emission()), two_sided=False if surface.double_sided is None else surface.double_sided, - beam_angle=surface.beam_angle, + beam_angle=surface.cutoff, ) self._scene.update_emission(emission_luisa) else: diff --git a/genesis/vis/visualizer.py b/genesis/vis/visualizer.py index a92ae846dd..86e82c763b 100644 --- a/genesis/vis/visualizer.py +++ b/genesis/vis/visualizer.py @@ -97,7 +97,7 @@ def __init__(self, scene, show_viewer, vis_options, viewer_options, renderer_opt if isinstance(renderer_options, gs.renderers.BatchRenderer): from .batch_renderer import BatchRenderer - self._renderer = self._batch_renderer = BatchRenderer(self, renderer_options) + self._renderer = self._batch_renderer = BatchRenderer(self, renderer_options, vis_options) elif isinstance(renderer_options, gs.renderers.RayTracer): from .raytracer import Raytracer @@ -130,7 +130,9 @@ def destroy(self): self.viewer_lock = None self._renderer = None - def add_camera(self, res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, env_idx, debug): + def add_camera( + self, res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise, near, far, env_idx, debug + ): cam_idx = len([camera for camera in self._cameras if camera.debug == debug]) camera = Camera( self, @@ -146,15 +148,23 @@ def add_camera(self, res, pos, lookat, up, model, fov, aperture, focus_dist, GUI GUI, spp, denoise, + near, + far, env_idx=env_idx, debug=debug, ) self._cameras.append(camera) return camera - def add_light(self, pos, dir, intensity, directional, castshadow, cutoff): + def add_mesh_light(self, mesh, color, intensity, pos, quat, revert_dir, double_sided, cutoff): + if self._raytracer is not None: + self._raytracer.add_mesh_light(mesh, color, intensity, pos, quat, revert_dir, double_sided, cutoff) + else: + gs.raise_exception("`add_mesh_light` is specific to raytracer renderer.") + + def add_light(self, pos, dir, color, intensity, directional, castshadow, cutoff, attenuation): if self._batch_renderer is not None: - self._batch_renderer.add_light(pos, dir, intensity, directional, castshadow, cutoff) + self._batch_renderer.add_light(pos, dir, color, intensity, directional, castshadow, cutoff, attenuation) else: gs.raise_exception("`add_light` is specific to batch renderer.") @@ -252,6 +262,12 @@ def update_visual_states(self): self._t = self._scene._t + def colorize_seg_idxc_arr(self, seg_idxc_arr): + if self._batch_renderer is not None: + return self._batch_renderer.colorize_seg_idxc_arr(seg_idxc_arr) + else: + return self._context.colorize_seg_idxc_arr(seg_idxc_arr) + # ------------------------------------------------------------------------------------ # ----------------------------------- properties ------------------------------------- # ------------------------------------------------------------------------------------ @@ -291,3 +307,10 @@ def has_display(self): @property def cameras(self): return self._cameras + + @property + def segmentation_idx_dict(self): + if self._batch_renderer is not None: + return self._batch_renderer.seg_idxc_map + else: + return self._context.seg_idxc_map diff --git a/pyproject.toml b/pyproject.toml index a3b291bb61..04c8ea67fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,10 @@ dependencies = [ "tetgen>=0.6.4", # Used for some advanced mesh processing such as `skeletonization` "PyGEL3D", - # Use to display camera images + # Used here and there to load and export images + # * 11.0 dramatically improves performance when exporting images to PNG + "Pillow>11.0", + # Use to display camera images and export images "opencv-python", # Use by `RigidGeom.visualize_sdf` to render SDF as level 0 marching cubes "scikit-image", @@ -57,7 +60,7 @@ dependencies = [ # Used for loading raytracing special texture images used by LuisaRender "OpenEXR", # Native batch renderer specifically designed for Genesis - "gs-madrona==0.0.3; platform_system=='Linux' and (platform_machine=='x86_64' or platform_machine=='AMD64')", + "gs-madrona==0.0.5; platform_system=='Linux' and (platform_machine=='x86_64' or platform_machine=='AMD64')", # Used for mesh simplification "fast_simplification>=0.1.12", # Surface reconstruction library for particle data from SPH simulations diff --git a/tests/test_render.py b/tests/test_render.py index 06c8542865..5919c35733 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -87,12 +87,6 @@ def test_render_api(show_viewer, renderer_type, renderer): rgb_arrs, depth_arrs, seg_arrs, normal_arrs = [], [], [], [] for rgb, depth, seg, normal in itertools.product((True, False), repeat=4): - if (seg or normal) and renderer_type in ( - RENDERER_TYPE.BATCHRENDER_RASTERIZER, - RENDERER_TYPE.BATCHRENDER_RAYTRACER, - ): - # Depth map and segmentation maps are not supported by Madrona for now. - continue rgb_arr, depth_arr, seg_arr, normal_arr = camera.render(rgb=rgb, depth=depth, segmentation=seg, normal=normal) if rgb: rgb_arrs.append(tensor_to_array(rgb_arr).astype(np.float32)) @@ -406,47 +400,52 @@ def test_render_api_advanced(tmp_path, n_envs, show_viewer, png_snapshot, render if IS_BATCHRENDER: # Note that the individual cameras is rendered alone first on purpose to make sure it works rgba_1, depth_1, seg_1, normal_1 = cam_1.render( - rgb=True, depth=True, segmentation=False, colorize_seg=False, normal=False + rgb=True, depth=True, segmentation=True, colorize_seg=True, normal=True ) rgba_all, depth_all, seg_all, normal_all = scene.render_all_cameras( - rgb=True, depth=True, segmentation=False, normal=False + rgb=True, depth=True, segmentation=True, colorize_seg=True, normal=True + ) + assert all(isinstance(img_data, torch.Tensor) for img_data in (rgba_1, depth_1, seg_1, normal_1)) + assert all( + isinstance(img_data, torch.Tensor) for img_data in (*rgba_all, *depth_all, *seg_all, *normal_all) ) - assert all(isinstance(img_data, torch.Tensor) for img_data in (rgba_1, depth_1)) - assert all(isinstance(img_data, torch.Tensor) for img_data in (*rgba_all, *depth_all)) else: # Emulate batch rendering which is not supported natively colorize_seg = False - rgba_all, depth_all, _, _ = zip( + rgba_all, depth_all, seg_all, normal_all = zip( *( - camera.render(rgb=True, depth=True, segmentation=False, normal=False) + camera.render(rgb=True, depth=True, segmentation=True, colorize_seg=True, normal=True) for camera in scene._visualizer._cameras if not camera.debug ) ) if n_envs > 0: - rgba_all, depth_all = ( + rgba_all, depth_all, seg_all, normal_all = ( tuple(np.swapaxes(np.stack(img_data, axis=0).reshape((n_envs, 3, *img_data[0].shape)), 0, 1)) - for img_data in (rgba_all, depth_all) + for img_data in (rgba_all, depth_all, seg_all, normal_all) ) - rgba_1, depth_1 = rgba_all[1], depth_all[1] + rgba_1, depth_1, seg_1, normal_1 = rgba_all[1], depth_all[1], seg_all[1], normal_all[1] # Check that the dimensions are valid batch_shape = (*((n_envs,) if n_envs else ()), *CAM_RES) assert len(rgba_all) == len(depth_all) == 3 - assert all(e.shape == (*batch_shape, 3) for e in (*rgba_all, rgba_1)) + assert all(e.shape == (*batch_shape, 3) for e in (*rgba_all, *seg_all, *normal_all, rgba_1, seg_1, normal_1)) assert all(e.shape == batch_shape for e in (*depth_all, depth_1)) # Check that the camera whose output was rendered individually is matching batched output for img_data_1, img_data_2 in ( (rgba_all[1], rgba_1), (depth_all[1], depth_1), + (seg_all[1], seg_1), + (normal_all[1], normal_1), ): assert_allclose(img_data_1, img_data_2, tol=gs.EPS) # Check that there is something to see here depth_normalized_all = tuple(as_grayscale_image(tensor_to_array(img_data)) for img_data in depth_all) frame_data = tuple( - tensor_to_array(img_data).astype(np.float32) for img_data in (*rgba_all, *depth_normalized_all) + tensor_to_array(img_data).astype(np.float32) + for img_data in (*rgba_all, *depth_normalized_all, *seg_all, *normal_all) ) for img_data in frame_data: for img_data_i in img_data if n_envs else (img_data,): @@ -454,8 +453,10 @@ def test_render_api_advanced(tmp_path, n_envs, show_viewer, png_snapshot, render # Export a few frames for later pixel-matching validation if i < 2: - exporter.export_frame_all_cameras(i, rgb=rgba_all, depth=depth_all) - exporter.export_frame_single_camera(i, cam_1.idx, rgb=rgba_1, depth=depth_1) + exporter.export_frame_all_cameras(i, rgb=rgba_all, depth=depth_all, segmentation=seg_all, normal=normal_all) + exporter.export_frame_single_camera( + i, cam_1.idx, rgb=rgba_1, depth=depth_1, segmentation=seg_1, normal=normal_1 + ) # Check that cameras are recording different part of the scene for rgb_diff in np.diff(frame_data[:3], axis=0): @@ -485,9 +486,9 @@ def test_render_api_advanced(tmp_path, n_envs, show_viewer, png_snapshot, render @pytest.mark.parametrize( "renderer_type", - [RENDERER_TYPE.RASTERIZER], + [RENDERER_TYPE.RASTERIZER, RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER], ) -@pytest.mark.parametrize("segmentation_level", ["entity", "link"]) +@pytest.mark.parametrize("segmentation_level", ["entity", "link", "geom"]) @pytest.mark.parametrize("particle_mode", ["visual", "particle"]) def test_segmentation_map(segmentation_level, particle_mode, renderer_type, renderer, show_viewer): """Test segmentation rendering.""" @@ -553,7 +554,7 @@ def test_segmentation_map(segmentation_level, particle_mode, renderer_type, rend scene.build() seg_num = len(materials) + (2 if segmentation_level == "entity" else 3) - idx_dict = camera.get_segmentation_idx_dict() + idx_dict = scene.segmentation_idx_dict assert len(idx_dict) == seg_num comp_key = 0 for seg_key in idx_dict.values(): @@ -564,15 +565,16 @@ def test_segmentation_map(segmentation_level, particle_mode, renderer_type, rend for i in range(2): scene.step() _, _, seg, _ = camera.render(rgb=False, depth=False, segmentation=True, colorize_seg=False, normal=False) + seg = tensor_to_array(seg) assert_array_equal(np.sort(np.unique(seg.flat)), np.arange(0, seg_num)) @pytest.mark.required @pytest.mark.parametrize( "renderer_type", - [RENDERER_TYPE.RASTERIZER], + [RENDERER_TYPE.RASTERIZER, RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER], ) -def test_point_cloud(renderer, show_viewer): +def test_point_cloud(renderer_type, renderer, show_viewer): N_ENVS = 2 CAM_RES = (256, 256) CAMERA_DIST = 8.0 @@ -580,11 +582,31 @@ def test_point_cloud(renderer, show_viewer): BOX_HALFSIZE = 1.0 SPHERE_RADIUS = 1.0 + IS_BATCHRENDER = renderer_type in (RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER) + BATCH_SHAPE = (N_ENVS,) if N_ENVS > 0 and IS_BATCHRENDER else () + scene = gs.Scene( renderer=renderer, show_viewer=show_viewer, show_FPS=False, ) + if renderer_type in (RENDERER_TYPE.BATCHRENDER_RASTERIZER, RENDERER_TYPE.BATCHRENDER_RAYTRACER): + scene.add_light( + pos=(0.0, 0.0, 1.5), + dir=(1.0, 1.0, -2.0), + directional=True, + castshadow=True, + cutoff=45.0, + intensity=0.5, + ) + scene.add_light( + pos=(4.0, -4.0, 4.0), + dir=(-1.0, 1.0, -1.0), + directional=False, + castshadow=True, + cutoff=45.0, + intensity=0.5, + ) scene.add_entity( morph=gs.morphs.Sphere( pos=(0.0, OBJ_OFFSET, 0.0), @@ -603,23 +625,26 @@ def test_point_cloud(renderer, show_viewer): res=CAM_RES, pos=(0.0, OBJ_OFFSET, CAMERA_DIST), lookat=(0.0, OBJ_OFFSET, 0.0), + near=2.0, + far=15.0, GUI=show_viewer, ) camera_box_1 = scene.add_camera( res=CAM_RES, pos=(0.0, -OBJ_OFFSET, CAMERA_DIST), lookat=(0.0, -OBJ_OFFSET, 0.0), + near=2.0, + far=15.0, GUI=show_viewer, ) camera_box_2 = scene.add_camera( res=CAM_RES, pos=np.array((CAMERA_DIST, CAMERA_DIST - OBJ_OFFSET, CAMERA_DIST)), lookat=(0.0, -OBJ_OFFSET, 0.0), + near=2.0, + far=15.0, GUI=show_viewer, ) - for camera in scene.visualizer.cameras: - camera._near = 2.0 - camera._far = 15.0 scene.build(n_envs=N_ENVS) if show_viewer: @@ -627,14 +652,14 @@ def test_point_cloud(renderer, show_viewer): camera.render(rgb=True, depth=True) point_cloud, mask = camera_box_1.render_pointcloud(world_frame=False) - assert point_cloud.shape == (*CAM_RES, 3) + assert point_cloud.shape == (*BATCH_SHAPE, *CAM_RES, 3) 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) - assert point_cloud.shape == (*CAM_RES, 3) + assert point_cloud.shape == (*BATCH_SHAPE, *CAM_RES, 3) 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)) @@ -643,14 +668,14 @@ def test_point_cloud(renderer, show_viewer): assert_allclose(np.linalg.norm(point_cloud, ord=float("inf"), axis=-1), BOX_HALFSIZE, atol=tol) point_cloud, mask = camera_box_2.render_pointcloud(world_frame=True) - assert point_cloud.shape == (*CAM_RES, 3) + assert point_cloud.shape == (*BATCH_SHAPE, *CAM_RES, 3) 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=tol) # It is not possible to get higher accuracy because of tesselation point_cloud, mask = camera_sphere.render_pointcloud(world_frame=False) - assert point_cloud.shape == (*CAM_RES, 3) + assert point_cloud.shape == (*BATCH_SHAPE, *CAM_RES, 3) 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) diff --git a/tests/utils.py b/tests/utils.py index 2b57431653..00ee309eeb 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -31,7 +31,7 @@ DEFAULT_BRANCH_NAME = "main" HUGGINGFACE_ASSETS_REVISION = "0c0bb46db0978a59524381194478cf390b3ff996" -HUGGINGFACE_SNAPSHOT_REVISION = "029b7d8932455ce31571d0660f89e28bcde4502f" +HUGGINGFACE_SNAPSHOT_REVISION = "9a192b8d4d34401b7e20d43601ff73f80516cb2b" MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS) IMAGE_EXTENSIONS = (".png", ".jpg")