Spaces:
Sleeping
Sleeping
| # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| ########################################################################### | |
| # OpenGL renderer example | |
| # | |
| # Demonstrates how to set up tiled rendering and retrieves the pixels from | |
| # OpenGLRenderer as a Warp array while keeping all memory on the GPU. | |
| # | |
| ########################################################################### | |
| import warp as wp | |
| import warp.render | |
| import numpy as np | |
| wp.init() | |
| # number of viewports to render in a single frame | |
| num_tiles = 4 | |
| # whether to split tiles into subplots | |
| split_up_tiles = True | |
| # whether to apply custom tile arrangement | |
| custom_tile_arrangement = False | |
| # whether to display the pixels in a matplotlib figure | |
| show_plot = True | |
| # whether to render depth image to a Warp array | |
| render_mode = "depth" | |
| renderer = wp.render.OpenGLRenderer(vsync=False) | |
| instance_ids = [] | |
| if custom_tile_arrangement: | |
| positions = [] | |
| sizes = [] | |
| else: | |
| positions = None | |
| sizes = None | |
| if num_tiles > 1: | |
| # set up instances to hide one of the capsules in each tile | |
| for i in range(num_tiles): | |
| instances = [j for j in np.arange(13) if j != i + 2] | |
| instance_ids.append(instances) | |
| if custom_tile_arrangement: | |
| angle = np.pi * 2.0 / num_tiles * i | |
| positions.append((int(np.cos(angle) * 150 + 250), int(np.sin(angle) * 150 + 250))) | |
| sizes.append((150, 150)) | |
| renderer.setup_tiled_rendering(instance_ids, tile_positions=positions, tile_sizes=sizes) | |
| renderer.render_ground() | |
| channels = 1 if render_mode == "depth" else 3 | |
| if show_plot: | |
| import matplotlib.pyplot as plt | |
| if split_up_tiles: | |
| pixels = wp.zeros((num_tiles, renderer.tile_height, renderer.tile_width, channels), dtype=wp.float32) | |
| ncols = int(np.ceil(np.sqrt(num_tiles))) | |
| nrows = int(np.ceil(num_tiles / float(ncols))) | |
| img_plots = [] | |
| aspect_ratio = renderer.tile_height / renderer.tile_width | |
| fig, axes = plt.subplots( | |
| ncols=ncols, | |
| nrows=nrows, | |
| constrained_layout=True, | |
| figsize=(ncols * 3.5, nrows * 3.5 * aspect_ratio), | |
| squeeze=False, | |
| sharex=True, | |
| sharey=True, | |
| num=1, | |
| ) | |
| tile_temp = np.zeros((renderer.tile_height, renderer.tile_width, channels), dtype=np.float32) | |
| for dim in range(ncols * nrows): | |
| ax = axes[dim // ncols, dim % ncols] | |
| if dim >= num_tiles: | |
| ax.axis("off") | |
| continue | |
| if render_mode == "depth": | |
| img_plots.append(ax.imshow(tile_temp, vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane)) | |
| else: | |
| img_plots.append(ax.imshow(tile_temp)) | |
| else: | |
| fig = plt.figure(1) | |
| pixels = wp.zeros((renderer.screen_height, renderer.screen_width, channels), dtype=wp.float32) | |
| if render_mode == "depth": | |
| img_plot = plt.imshow(pixels.numpy(), vmin=renderer.camera_near_plane, vmax=renderer.camera_far_plane) | |
| else: | |
| img_plot = plt.imshow(pixels.numpy()) | |
| plt.ion() | |
| plt.show() | |
| while renderer.is_running(): | |
| time = renderer.clock_time | |
| renderer.begin_frame(time) | |
| for i in range(10): | |
| renderer.render_capsule( | |
| f"capsule_{i}", [i - 5.0, np.sin(time + i * 0.2), -3.0], [0.0, 0.0, 0.0, 1.0], radius=0.5, half_height=0.8 | |
| ) | |
| renderer.render_cylinder( | |
| "cylinder", | |
| [3.2, 1.0, np.sin(time + 0.5)], | |
| np.array(wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.sin(time + 0.5))), | |
| radius=0.5, | |
| half_height=0.8, | |
| ) | |
| renderer.render_cone( | |
| "cone", | |
| [-1.2, 1.0, 0.0], | |
| np.array(wp.quat_from_axis_angle(wp.vec3(0.707, 0.707, 0.0), time)), | |
| radius=0.5, | |
| half_height=0.8, | |
| ) | |
| renderer.end_frame() | |
| if show_plot and plt.fignum_exists(1): | |
| if split_up_tiles: | |
| pixel_shape = (num_tiles, renderer.tile_height, renderer.tile_width, channels) | |
| else: | |
| pixel_shape = (renderer.screen_height, renderer.screen_width, channels) | |
| if pixel_shape != pixels.shape: | |
| # make sure we resize the pixels array to the right dimensions if the user resizes the window | |
| pixels = wp.zeros(pixel_shape, dtype=wp.float32) | |
| renderer.get_pixels(pixels, split_up_tiles=split_up_tiles, mode=render_mode) | |
| if split_up_tiles: | |
| pixels_np = pixels.numpy() | |
| for i, img_plot in enumerate(img_plots): | |
| img_plot.set_data(pixels_np[i]) | |
| else: | |
| img_plot.set_data(pixels.numpy()) | |
| fig.canvas.draw() | |
| fig.canvas.flush_events() | |
| renderer.clear() | |