[FEATURE] Add New Features for Madrona Batch Rendering (#1563)

Co-authored-by: Alexis DUBURCQ <alexis.duburcq@gmail.com>

Author: ACMLCZH

examples/rigid/single_franka_batch_render.py

@@ -13,16 +13,20 @@ 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 ##########################
     gs.init(backend=gs.cpu if args.cpu else gs.gpu)
 
     ########################## 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")
 

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

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

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.