Spaces:
Sleeping
Sleeping
File size: 5,085 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | # 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()
|