batch render render array of cam

examples/perf_benchmark/benchmark_madrona.py

@@ -30,7 +30,6 @@ def init_gs(benchmark_args):
             ),
         renderer = gs.options.renderers.BatchRenderer(
             use_rasterizer=benchmark_args.rasterizer,
-            batch_render_res=(benchmark_args.resX, benchmark_args.resY),
         )            
     )
 
@@ -45,6 +44,7 @@ def init_gs(benchmark_args):
 
     ########################## cameras ##########################
     cam_0 = scene.add_camera(
+        res=(benchmark_args.resX, benchmark_args.resY),
         pos=(benchmark_args.camera_posX, benchmark_args.camera_posY, benchmark_args.camera_posZ),
         lookat=(benchmark_args.camera_lookatX, benchmark_args.camera_lookatY, benchmark_args.camera_lookatZ),
         fov=benchmark_args.camera_fov,

examples/perf_benchmark/benchmark_profiler.py

@@ -5,7 +5,7 @@
 class BenchmarkProfiler:
     def __init__(self, n_steps, n_envs):
         self.reset(n_steps)
-        self.n_envs = n_envs
+        self.n_envs = max(n_envs, 1)
 
     def reset(self, n_steps):
         self.n_steps = n_steps

examples/rigid/batch_render_with_ppo.py

@@ -63,7 +63,6 @@ def __init__(self, num_envs: int = 32, res: tuple[int, int] = (64, 64), max_step
             ),
             renderer=gs.options.renderers.BatchRenderer(
                 use_rasterizer=False,
-                batch_render_res=res,
             ),
         )
 
@@ -87,7 +86,7 @@ def __init__(self, num_envs: int = 32, res: tuple[int, int] = (64, 64), max_step
         # hand_cam.attach(self.robot.links[6], trans_to_T(np.array([0.0, 0.5, 0.0])))
 
         # overview cam
-        self.scene.add_camera(pos=(1.5, 0.0, 1.5), lookat=(0.0, 0.0, 0.5), fov=45, GUI=False)
+        self.scene.add_camera(res=res, pos=(1.5, 0.0, 1.5), lookat=(0.0, 0.0, 0.5), fov=45, GUI=False)
 
         self.scene.build(n_envs=num_envs)
 

examples/rigid/single_franka_batch_render.py

@@ -30,7 +30,6 @@ def main():
         ),
         renderer = gs.options.renderers.BatchRenderer(
             use_rasterizer=False,
-            batch_render_res=(512, 512),
         )
     )
 
@@ -45,13 +44,15 @@ def main():
 
     ########################## cameras ##########################
     cam_0 = scene.add_camera(
+        res=(512, 512),
         pos=(1.5, 0.5, 1.5),
         lookat=(0.0, 0.0, 0.5),
         fov=45,
         GUI=True,
     )
     cam_0.attach(franka.links[6], trans_to_T(np.array([0.0, 0.5, 0.0])))
     cam_1 = scene.add_camera(
+        res=(512, 512),
         pos=(1.5, -0.5, 1.5),
         lookat=(0.0, 0.0, 0.5),
         fov=45,
@@ -79,18 +80,10 @@ def main():
     is_render_all_cams = True
     scene.build(n_envs=n_envs)
 
-    # Warmup
-    scene.step()
-    rgb, depth, _, _ = scene.render_all_cams()
-
     # Create an image exporter
     output_dir = 'img_output/test'
     exporter = FrameImageExporter(output_dir)
 
-    # Timer
-    from time import time
-    start_time = time()
-
     for i in range(n_steps):
         scene.step()
         if is_render_all_cams:
@@ -99,14 +92,6 @@ def main():
         else:
             rgb, depth, _, _ = cam_0.render()
             exporter.export_frame_single_cam(i, cam_0.idx, rgb=rgb, depth=depth)
-
-    end_time = time()
-    actual_n_envs = n_envs if n_envs > 0 else 1
-    print(f'n_envs: {n_envs}')
-    print(f'Time taken: {end_time - start_time} seconds')
-    print(f'Time taken per env: {(end_time - start_time) / actual_n_envs} seconds')
-    print(f'FPS: {actual_n_envs * n_steps / (end_time - start_time)}')
-    print(f'FPS per env: {n_steps / (end_time - start_time)}')
 
 
 if __name__ == "__main__":

genesis/engine/scene.py

@@ -1068,6 +1068,10 @@ def draw_debug_path(self, qposs, entity, link_idx=-1, density=0.3, frame_scaling
     def render_all_cams(self, force_render=False):
         """
         Render the scene for all cameras using the batch renderer.
+
+        Returns:
+            A list of tensors of shape (n_envs, H, W, 3) if rgb is not None, otherwise a list of tensors of shape (n_envs, H, W, 1) if depth is not None.
+            If n_envs ==0, the first dimension of the tensor is squeezed.
         """
         return self._visualizer.batch_renderer.render(force_render=force_render)
 

genesis/options/renderers.py

@@ -109,9 +109,6 @@ class BatchRenderer(RendererOptions):
     ----------
     use_rasterizer : bool, optional
         Whether to use the rasterizer renderer. Defaults to False.
-    batch_render_res : tuple, optional
-        The resolution of the batch render. Defaults to (128, 128).
     """
 
-    use_rasterizer: bool = False
-    batch_render_res: tuple = (128, 128)
+    use_rasterizer: bool = False

genesis/utils/image_exporter.py

@@ -6,22 +6,22 @@
 
 class FrameImageExporter:
     @staticmethod
-    def _export_frame_rgb_cam(export_dir, i_env, i_cam, camera_name, i_step, rgb):
-        rgb = rgb[i_env, i_cam, ..., [2, 1, 0]].cpu().numpy()
-        cv2.imwrite(f'{export_dir}/rgb_env{i_env}_{camera_name}_{i_step:03d}.png', rgb)
+    def _export_frame_rgb_cam(export_dir, i_cam, i_env, i_step, rgb):
+        rgb = rgb[i_env, ..., [2, 1, 0]].cpu().numpy()
+        cv2.imwrite(f'{export_dir}/rgb_cam{i_cam}_env{i_env}_{i_step:03d}.png', rgb)
 
     @staticmethod
-    def _export_frame_depth_cam(export_dir, i_env, i_cam, camera_name, i_step, depth):
-        depth = depth[i_env, i_cam].cpu().numpy()
-        cv2.imwrite(f'{export_dir}/depth_env{i_env}_{camera_name}_{i_step:03d}.png', depth)
+    def _export_frame_depth_cam(export_dir, i_cam, i_env, i_step, depth):
+        depth = depth[i_env].cpu().numpy()
+        cv2.imwrite(f'{export_dir}/depth_cam{i_cam}_env{i_env}_{i_step:03d}.png', depth)
 
     @staticmethod
     def _worker_export_frame_cam(args):
-        export_dir, i_env, i_cam, camera_name, rgb, depth, i_step = args
+        export_dir, i_cam, i_env, rgb, depth, i_step = args
         if rgb is not None:
-            FrameImageExporter._export_frame_rgb_cam(export_dir, i_env, i_cam, camera_name, i_step, rgb)
+            FrameImageExporter._export_frame_rgb_cam(export_dir, i_cam, i_env, i_step, rgb)
         if depth is not None:
-            FrameImageExporter._export_frame_depth_cam(export_dir, i_env, i_cam, camera_name, i_step, depth)
+            FrameImageExporter._export_frame_depth_cam(export_dir, i_cam, i_env, i_step, depth)
 
     def __init__(self, export_dir, depth_clip_max=100, depth_scale='log'):
         self.export_dir = export_dir
@@ -52,9 +52,6 @@ def _normalize_depth(self, depth):
 
         # Normalize to 0-255 range
         return ((depth - depth_min) / (depth_max - depth_min) * 255).to(torch.uint8)
-
-    def _get_camera_name(self, i_cam):
-        return 'cam' + str(i_cam)
 
     def export_frame_all_cams(self, i_step, camera_idx=None, rgb=None, depth=None):
         """
@@ -63,72 +60,57 @@ def export_frame_all_cams(self, i_step, camera_idx=None, rgb=None, depth=None):
         Args:
             i_step: The current step index.
             camera_idx: array of indices of cameras to export. If None, all cameras are exported.
-            rgb: RGB image tensor of shape (n_envs, n_cams, H, W, 3).
-            depth: Depth tensor of shape (n_envs, n_cams, H, W).
+            rgb: RGB image is a list of tensors of shape (n_envs, H, W, 3).
+            depth: Depth image is a list of tensors of shape (n_envs, H, W).
         """
         if rgb is None and depth is None:
             print("No rgb or depth to export")
             return
-
         if rgb is not None:
-            if rgb.ndim == 4:
-                rgb = rgb.unsqueeze(0)
-            assert rgb.ndim == 5, "rgb must be of shape (n_envs, n_cams, H, W, 3)"
+            assert isinstance(rgb, list) and len(rgb) > 0, "rgb must be a list of tensors with length > 0"
         if depth is not None:
-            if depth.ndim == 4:
-                depth = depth.unsqueeze(0)
-            depth = self._normalize_depth(depth)
-            assert depth.ndim == 5, "depth must be of shape (n_envs, n_cams, H, W, 1)"
-
+            assert isinstance(depth, list) and len(depth) > 0, "depth must be a list of tensors with length > 0"
         if camera_idx is None:
-            camera_idx = range(rgb.shape[1]) if rgb is not None else range(depth.shape[1])
-        env_idx = range(rgb.shape[0]) if rgb is not None else range(depth.shape[0])
-        args_list = [(self.export_dir, i_env, i_cam, self._get_camera_name(i_cam), rgb, depth, i_step) 
-                     for i_env in env_idx for i_cam in camera_idx]
-        with ThreadPoolExecutor() as executor:
-            executor.map(FrameImageExporter._worker_export_frame_cam, args_list)
+            camera_idx = range(len(rgb)) if rgb is not None else range(len(depth))
+        for i_cam in camera_idx:
+            rgb_cam = rgb[i_cam] if i_cam < len(rgb) else None
+            depth_cam = depth[i_cam] if i_cam < len(depth) else None
+            if rgb_cam is not None or depth_cam is not None:
+                self.export_frame_single_cam(i_step, i_cam, rgb_cam, depth_cam)
 
     def export_frame_single_cam(self, i_step, i_cam, rgb=None, depth=None):
         """
         Export frames for a single camera.
 
         Args:
             i_step: The current step index.
-            cam_idx: The index of the camera.
+            i_cam: The index of the camera.
             rgb: RGB image tensor of shape (n_envs, H, W, 3).
             depth: Depth tensor of shape (n_envs, H, W).
         """
         if rgb is not None:
             if isinstance(rgb, np.ndarray):
                 rgb = torch.from_numpy(rgb.copy())
 
-            # Unsqueeze rgb to (n_envs, 1, H, W, 3) if n_envs > 0
-            if rgb.ndim == 4:
-                rgb = rgb.unsqueeze(1)
-            elif rgb.ndim == 3:
-                rgb = rgb.unsqueeze(0).unsqueeze(0)
-            else:
-                raise ValueError(f"Invalid rgb shape: {rgb.shape}")
-            assert rgb.ndim == 5, "rgb must be of shape (n_envs, H, W, 3)"
+            # Unsqueeze rgb to (n_envs, H, W, 3) if n_envs > 0
+            if rgb.ndim == 3:
+                rgb = rgb.unsqueeze(0)
+            assert rgb.ndim == 4, "rgb must be of shape (n_envs, H, W, 3)"
 
         if depth is not None:
             if isinstance(depth, np.ndarray):
                 depth = torch.from_numpy(depth.copy())
 
-            # Unsqueeze depth to (n_envs, 1, H, W, 1) if n_envs > 0
-            if depth.ndim == 4:
-                depth = depth.unsqueeze(1)
-            elif depth.ndim == 3:
-                depth = depth.unsqueeze(0).unsqueeze(0)
+            # Unsqueeze depth to (n_envs, H, W, 1) if n_envs > 0
+            if depth.ndim == 3:
+                depth = depth.unsqueeze(0)
             elif depth.ndim == 2:
-                depth = depth.unsqueeze(0).unsqueeze(0).unsqueeze(4)
-            else:
-                raise ValueError(f"Invalid depth shape: {depth.shape}")                    
+                depth = depth.unsqueeze(0).unsqueeze(3)          
             depth = self._normalize_depth(depth)
-            assert depth.ndim == 5, "depth must be of shape (n_envs, H, W, 1)"
+            assert depth.ndim == 4, "depth must be of shape (n_envs, H, W, 1)"
 
         env_idx = range(rgb.shape[0]) if rgb is not None else range(depth.shape[0])
-        args_list = [(self.export_dir, i_env, 0, self._get_camera_name(i_cam), rgb, depth, i_step) 
+        args_list = [(self.export_dir, i_cam, i_env, rgb, depth, i_step) 
                      for i_env in env_idx]
         with ThreadPoolExecutor() as executor:
             executor.map(FrameImageExporter._worker_export_frame_cam, args_list) 

genesis/vis/batch_renderer.py

@@ -79,7 +79,7 @@ def build(self):
         rigid = self._visualizer.scene.rigid_solver
         device = torch.cuda.current_device()
         n_envs = self._visualizer.scene.n_envs if self._visualizer.scene.n_envs > 0 else 1
-        res = self._renderer_options.batch_render_res
+        res = cameras[0].res
         use_rasterizer = self._renderer_options.use_rasterizer
 
         # Cameras
@@ -129,6 +129,11 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_
         """
         Render all cameras in the batch.
         """
+        if(normal):
+            raise NotImplementedError("Normal rendering is not implemented")
+        if(segmentation):
+            raise NotImplementedError("Segmentation rendering is not implemented")
+
         if(not force_render and self._last_t == self._visualizer.scene.t):
             return self._rgb_torch, self._depth_torch, None, None
 
@@ -148,8 +153,18 @@ def render(self, rgb=True, depth=False, segmentation=False, normal=False, force_
 
         # Squeeze the first dimension of the output if n_envs == 0
         if self._visualizer.scene.n_envs == 0:
-            self._rgb_torch = self._rgb_torch.squeeze(0)
-            self._depth_torch = self._depth_torch.squeeze(0)
+            if(self._rgb_torch.ndim == 4):
+                self._rgb_torch = self._rgb_torch.squeeze(0)
+            if(self._depth_torch.ndim == 4):
+                self._depth_torch = self._depth_torch.squeeze(0)
+
+        # swap the first two dimensions of the output
+        self._rgb_torch = self._rgb_torch.swapaxes(0, 1)
+        self._depth_torch = self._depth_torch.swapaxes(0, 1)
+
+        # Create a list of tensors pointing to each sub tensor
+        self._rgb_torch = [self._rgb_torch[i] for i in range(self._rgb_torch.shape[0])]
+        self._depth_torch = [self._depth_torch[i] for i in range(self._depth_torch.shape[0])]
         return self._rgb_torch, self._depth_torch, None, None
 
     def destroy(self):

genesis/vis/camera.py

@@ -186,17 +186,18 @@ def _batch_render(self, rgb=True, depth=False, segmentation=False, colorize_seg=
         assert self._visualizer._use_batch_renderer, "Batch renderer is not enabled."
 
         rgb_arr, depth_arr, seg_arr, normal_arr = self._batch_renderer.render(rgb, depth)
-        # If n_envs > 0, the first dimension of the output is env. The second dimension of the array is camera.
-        # If n_envs == 0, the first dimension of the output is camera.
+        # The first dimension of the array is camera.
+        # If n_envs > 0, the second dimension of the output is env. 
+        # If n_envs == 0, the second dimension of the output is camera.
         # Only return the current camera's image
         if rgb_arr is not None:
-            rgb_arr = rgb_arr[:, self._idx] if rgb_arr.ndim == 5 else rgb_arr[self._idx]
+            rgb_arr = rgb_arr[self._idx]
         if depth_arr is not None:
-            depth_arr = depth_arr[:, self._idx] if depth_arr.ndim == 5 else depth_arr[self._idx]
+            depth_arr = depth_arr[self._idx]
         if seg_arr is not None:
-            seg_arr = seg_arr[:, self._idx] if seg_arr.ndim == 5 else seg_arr[self._idx]
+            seg_arr = seg_arr[self._idx]
         if normal_arr is not None:
-            normal_arr = normal_arr[:, self._idx] if normal_arr.ndim == 5 else normal_arr[self._idx]
+            normal_arr = normal_arr[self._idx]
         return rgb_arr, depth_arr, seg_arr, normal_arr
 
     @gs.assert_built

genesis/vis/visualizer.py

@@ -89,7 +89,6 @@ def __init__(self, scene, show_viewer, vis_options, viewer_options, renderer_opt
             self._batch_renderer = BatchRenderer(self, renderer_options)
             self._renderer = self._batch_renderer
             self._raytracer = None
-            self.batch_camera_res = renderer_options.batch_render_res
         elif isinstance(renderer_options, gs.renderers.RayTracer):
             from .raytracer import Raytracer
             self._renderer = self._raytracer = Raytracer(renderer_options, vis_options)
@@ -122,10 +121,6 @@ def destroy(self):
         self._renderer = None
 
     def add_camera(self, res, pos, lookat, up, model, fov, aperture, focus_dist, GUI, spp, denoise):
-        if(self._use_batch_renderer):
-            if(res != self.batch_camera_res):
-                gs.logger.warning("Camera resolution mismatch with batch renderer resolution. Overwriting camera resolution.")
-                res = self.batch_camera_res
         camera = Camera(
             self, len(self._cameras), model, res, pos, lookat, up, fov, aperture, focus_dist, GUI, spp, denoise
         )