From 3c39f8c2528853c9c26593e895fabf5063209ef5 Mon Sep 17 00:00:00 2001 From: Hongyi Yu Date: Wed, 18 Jun 2025 09:31:44 -0700 Subject: [PATCH 1/4] batch render render array of cam --- examples/perf_benchmark/benchmark_profiler.py | 2 +- examples/rigid/single_franka_batch_render.py | 16 ---- genesis/utils/image_exporter.py | 78 +++++++------------ genesis/vis/batch_renderer.py | 19 ++++- genesis/vis/camera.py | 13 ++-- 5 files changed, 55 insertions(+), 73 deletions(-) diff --git a/examples/perf_benchmark/benchmark_profiler.py b/examples/perf_benchmark/benchmark_profiler.py index 6a245cd82..375cea6aa 100644 --- a/examples/perf_benchmark/benchmark_profiler.py +++ b/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 diff --git a/examples/rigid/single_franka_batch_render.py b/examples/rigid/single_franka_batch_render.py index 774df8354..43a6c1cde 100644 --- a/examples/rigid/single_franka_batch_render.py +++ b/examples/rigid/single_franka_batch_render.py @@ -79,18 +79,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 +91,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__": diff --git a/genesis/utils/image_exporter.py b/genesis/utils/image_exporter.py index fe8c6382b..7fdd4deb2 100644 --- a/genesis/utils/image_exporter.py +++ b/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,30 +60,23 @@ 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): """ @@ -94,7 +84,7 @@ def export_frame_single_cam(self, i_step, i_cam, rgb=None, depth=None): 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). """ @@ -102,33 +92,25 @@ def export_frame_single_cam(self, i_step, i_cam, rgb=None, depth=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) diff --git a/genesis/vis/batch_renderer.py b/genesis/vis/batch_renderer.py index abe8035bb..39285ce77 100644 --- a/genesis/vis/batch_renderer.py +++ b/genesis/vis/batch_renderer.py @@ -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): diff --git a/genesis/vis/camera.py b/genesis/vis/camera.py index f288610ee..9f5c69da4 100644 --- a/genesis/vis/camera.py +++ b/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 From b3e5d0cb15046fae47b6320d5b452f61e2e3d219 Mon Sep 17 00:00:00 2001 From: Hongyi Yu Date: Wed, 18 Jun 2025 09:43:08 -0700 Subject: [PATCH 2/4] Use camera[0].res as batch render res --- examples/perf_benchmark/benchmark_madrona.py | 2 +- examples/rigid/batch_render_with_ppo.py | 3 +-- examples/rigid/single_franka_batch_render.py | 1 - genesis/options/renderers.py | 5 +---- genesis/vis/batch_renderer.py | 2 +- genesis/vis/visualizer.py | 5 ----- 6 files changed, 4 insertions(+), 14 deletions(-) diff --git a/examples/perf_benchmark/benchmark_madrona.py b/examples/perf_benchmark/benchmark_madrona.py index 90284ade9..7f6c0dc2d 100644 --- a/examples/perf_benchmark/benchmark_madrona.py +++ b/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, diff --git a/examples/rigid/batch_render_with_ppo.py b/examples/rigid/batch_render_with_ppo.py index 4fc84b20f..d3a80618b 100644 --- a/examples/rigid/batch_render_with_ppo.py +++ b/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) diff --git a/examples/rigid/single_franka_batch_render.py b/examples/rigid/single_franka_batch_render.py index 43a6c1cde..541d20469 100644 --- a/examples/rigid/single_franka_batch_render.py +++ b/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), ) ) diff --git a/genesis/options/renderers.py b/genesis/options/renderers.py index 86158df0b..cabab4b25 100644 --- a/genesis/options/renderers.py +++ b/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) \ No newline at end of file + use_rasterizer: bool = False \ No newline at end of file diff --git a/genesis/vis/batch_renderer.py b/genesis/vis/batch_renderer.py index 39285ce77..303fea5db 100644 --- a/genesis/vis/batch_renderer.py +++ b/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 diff --git a/genesis/vis/visualizer.py b/genesis/vis/visualizer.py index 03cb3f801..fa34c4608 100644 --- a/genesis/vis/visualizer.py +++ b/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 ) From 197b236bf0cb9124eb1ead1cab18291057c88050 Mon Sep 17 00:00:00 2001 From: Hongyi Yu Date: Wed, 18 Jun 2025 09:46:45 -0700 Subject: [PATCH 3/4] Batch render res --- examples/rigid/single_franka_batch_render.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/rigid/single_franka_batch_render.py b/examples/rigid/single_franka_batch_render.py index 541d20469..5ecb2fb46 100644 --- a/examples/rigid/single_franka_batch_render.py +++ b/examples/rigid/single_franka_batch_render.py @@ -44,6 +44,7 @@ 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, @@ -51,6 +52,7 @@ def main(): ) 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, From 556b32628fb4d0a70906dd9ff218e988d306bcc9 Mon Sep 17 00:00:00 2001 From: Hongyi Yu Date: Wed, 18 Jun 2025 10:06:45 -0700 Subject: [PATCH 4/4] Update comments --- genesis/engine/scene.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/genesis/engine/scene.py b/genesis/engine/scene.py index 537c1b7ef..f9620fff8 100644 --- a/genesis/engine/scene.py +++ b/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)