Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions genesis/engine/scene.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
import pickle
import sys
import time

import numpy as np
Expand Down Expand Up @@ -526,7 +527,7 @@ def add_camera(
focus_dist=None,
GUI=False,
spp=256,
denoise=True,
denoise=None,
env_idx=None,
):
"""
Expand Down Expand Up @@ -564,14 +565,16 @@ def add_camera(
Samples per pixel. Only available when using RayTracer renderer. Defaults to 256.
denoise : bool
Whether to denoise the camera's rendered image. Only available when using the RayTracer renderer. Defaults
to True. If OptiX denoiser is not available in your platform, consider enabling the OIDN denoiser option
when building the RayTracer.
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.
Returns
-------
camera : genesis.Camera
The created camera object.
"""
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
)
Expand Down
29 changes: 22 additions & 7 deletions genesis/ext/pyrender/offscreen.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,16 +173,22 @@ def render(
else:
if flags & RenderFlags.ENV_SEPARATE:
gs.raise_exception("'env_separate_rigid=True' not supported on this platform.")
if normal:
gs.raise_exception("'normal=True' not supported on this platform.")
renderer.render(scene, flags, seg_node_map)
glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
glReadBuffer(GL_FRONT)
if depth:
depth = renderer.read_depth_buf()
z_near = scene.main_camera_node.camera.znear
z_far = scene.main_camera_node.camera.zfar
if z_far is None:
z_far = -1.0
depth_arr = renderer.jit.read_depth_buf(self.viewport_height, self.viewport_width, z_near, z_far)
depth_arr = renderer._resize_image(depth_arr, antialias=not seg)
if flags & RenderFlags.DEPTH_ONLY:
retval = (depth,)
retval = (depth_arr,)
else:
color = renderer.read_color_buf()
retval = (color, depth)
color_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False)
color_arr = renderer._resize_image(color_arr, antialias=not seg)
retval = (color_arr, depth_arr) if depth else (color_arr,)
else:
retval = ()

Expand All @@ -207,7 +213,16 @@ def get_program(self, vertex_shader, fragment_shader, geometry_shader=None, defi
if env_separate_rigid:
flags |= RenderFlags.ENV_SEPARATE
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)

if self._platform.supports_framebuffers():
normal_arr, *_ = renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
else:
glBindFramebuffer(GL_READ_FRAMEBUFFER, 0)
glReadBuffer(GL_FRONT)
renderer.render(scene, flags, is_first_pass=False, force_skip_shadows=True)
normal_arr = renderer.jit.read_color_buf(self.viewport_height, self.viewport_width, rgba=False)
normal_arr = renderer._resize_image(normal_arr, antialias=not seg)

retval = (*retval, normal_arr)

renderer._program_cache = old_cache
Expand Down
5 changes: 3 additions & 2 deletions genesis/ext/pyrender/viewer.py
Original file line number Diff line number Diff line change
Expand Up @@ -1063,11 +1063,12 @@ def _save_image(self):
filename = self._get_save_filename(["png", "jpg", "gif", "all"])
if filename is not None:
self.viewer_flags["save_directory"] = os.path.dirname(filename)
imageio.imwrite(filename, self._renderer.read_color_buf())
data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False)
imageio.imwrite(filename, img_arr)

def _record(self):
"""Save another frame for the GIF."""
data = self._renderer.read_color_buf()
data = self._renderer.jit.read_color_buf(*self._viewport_size, rgba=False)
if not np.all(data == 0.0):
self.video_recorder.write_frame(data)

Expand Down
40 changes: 15 additions & 25 deletions genesis/vis/camera.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,13 @@ def __init__(
self._focus_dist = np.linalg.norm(np.asarray(lookat) - np.asarray(pos))

def build(self):
n_envs = max(self._visualizer.scene.n_envs, 1)
self._multi_env_pos_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
self._multi_env_lookat_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
self._multi_env_up_tensor = torch.empty((n_envs, 3), dtype=gs.tc_float, device=gs.device)
self._multi_env_transform_tensor = torch.empty((n_envs, 4, 4), dtype=gs.tc_float, device=gs.device)
self._multi_env_quat_tensor = torch.empty((n_envs, 4), dtype=gs.tc_float, device=gs.device)

self._envs_offset = torch.as_tensor(self._visualizer._scene.envs_offset, dtype=gs.tc_float, device=gs.device)

self._batch_renderer = self._visualizer.batch_renderer
Expand All @@ -154,7 +161,12 @@ def build(self):
self._other_stacked = self._visualizer._context.env_separate_rigid

self._is_built = True
self.setup_initial_env_poses()
self.set_pose(
transform=self._initial_transform, pos=self._initial_pos, lookat=self._initial_lookat, up=self._initial_up
)
# FIXME: For some reason, it is necessary to update the camera twice...
if self._raytracer is not None:
self._raytracer.update_camera(self)

def attach(self, rigid_link, offset_T):
"""
Expand Down Expand Up @@ -296,8 +308,8 @@ def _batch_render(
depth=False,
segmentation=False,
normal=False,
antialiasing=True,
force_render=False,
antialiasing=False,
):
"""
Render the camera view with batch renderer.
Expand Down Expand Up @@ -371,7 +383,7 @@ def render(
normal_arr : np.ndarray
The rendered surface normal(s).
"""
rgb_arr, depth_arr, seg_arr, seg_color_arr, normal_arr = None, None, None, None, None
rgb_arr, depth_arr, seg_arr, seg_color_arr, seg_idxc_arr, normal_arr = None, None, None, None, None, None

if self._batch_renderer is not None:
rgb_arr, depth_arr, seg_idxc_arr, normal_arr = self._batch_render(
Expand Down Expand Up @@ -521,28 +533,6 @@ def render_pointcloud(self, world_frame=True):
point_cloud = point_cloud[:, :3].reshape((*depth_arr.shape, 3))
return point_cloud, mask

@gs.assert_built
def setup_initial_env_poses(self):
"""
Setup the camera poses for multiple environments.
"""
if self._initial_transform is not None:
assert self._initial_transform.shape == (4, 4)
self._initial_pos, self._initial_lookat, self._initial_up = gu.T_to_pos_lookat_up(self._initial_transform)
else:
self._initial_transform = gu.pos_lookat_up_to_T(self._initial_pos, self._initial_lookat, self._initial_up)

n_envs = max(self._visualizer.scene.n_envs, 1)
self._multi_env_pos_tensor = self._initial_pos.repeat((n_envs, 1))
self._multi_env_lookat_tensor = self._initial_lookat.repeat((n_envs, 1))
self._multi_env_up_tensor = self._initial_up.repeat((n_envs, 1))
self._multi_env_transform_tensor = self._initial_transform.repeat((n_envs, 1, 1))
self._multi_env_quat_tensor = _T_to_quat_for_madrona(self._multi_env_transform_tensor)

self._rasterizer.update_camera(self)
if self._raytracer is not None:
self._raytracer.update_camera(self)

@gs.assert_built
def set_pose(self, transform=None, pos=None, lookat=None, up=None, env_idx=None):
"""
Expand Down
61 changes: 31 additions & 30 deletions genesis/vis/rasterizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,36 +69,37 @@ def render_camera(self, camera, rgb=True, depth=False, segmentation=False, norma
self._context.buffer.clear()

# Render
if rgb or depth or normal:
retval = self._renderer.render(
self._context._scene,
self._camera_targets[camera.uid],
camera_node=self._camera_nodes[camera.uid],
env_separate_rigid=self._context.env_separate_rigid,
rgb=rgb,
normal=normal,
seg=False,
depth=depth,
plane_reflection=rgb and self._context.plane_reflection,
shadow=rgb and self._context.shadow,
)

if segmentation:
seg_idxc_rgb_arr, *_ = self._renderer.render(
self._context._scene,
self._camera_targets[camera.uid],
camera_node=self._camera_nodes[camera.uid],
env_separate_rigid=self._context.env_separate_rigid,
rgb=False,
normal=False,
seg=True,
depth=False,
plane_reflection=False,
shadow=False,
)

# Unset the context
self._renderer.make_uncurrent()
try:
if rgb or depth or normal:
retval = self._renderer.render(
self._context._scene,
self._camera_targets[camera.uid],
camera_node=self._camera_nodes[camera.uid],
env_separate_rigid=self._context.env_separate_rigid,
rgb=rgb,
normal=normal,
seg=False,
depth=depth,
plane_reflection=rgb and self._context.plane_reflection,
shadow=rgb and self._context.shadow,
)

if segmentation:
seg_idxc_rgb_arr, *_ = self._renderer.render(
self._context._scene,
self._camera_targets[camera.uid],
camera_node=self._camera_nodes[camera.uid],
env_separate_rigid=self._context.env_separate_rigid,
rgb=False,
normal=False,
seg=True,
depth=False,
plane_reflection=False,
shadow=False,
)
finally:
# Unset the context
self._renderer.make_uncurrent()
else:
# Render
if rgb or depth or normal:
Expand Down
4 changes: 2 additions & 2 deletions genesis/vis/raytracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -604,7 +604,7 @@ def add_camera(self, camera):
if camera_model == "pinhole":
self._cameras[camera_name] = LuisaRenderPy.PinholeCamera(
name=camera_name,
pose=self.get_transform(camera.transform),
pose=self.get_transform(np.eye(4)),
film=LuisaRenderPy.Film(resolution=camera.res),
filter=LuisaRenderPy.Filter(),
spp=camera.spp,
Expand All @@ -613,7 +613,7 @@ def add_camera(self, camera):
elif camera_model == "thinlens":
self._cameras[camera_name] = LuisaRenderPy.ThinLensCamera(
name=camera_name,
pose=self.get_transform(camera.transform),
pose=self.get_transform(np.eye(4)),
film=LuisaRenderPy.Film(resolution=camera.res),
filter=LuisaRenderPy.Filter(),
spp=camera.spp,
Expand Down
4 changes: 3 additions & 1 deletion tests/test_render.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def test_segmentation(segmentation_level, particle_mode):

@pytest.mark.required
@pytest.mark.flaky(reruns=3, condition=(sys.platform == "darwin"))
def test_batched_offscreen_rendering(show_viewer, tol):
def test_batched_offscreen_rendering(tmp_path, show_viewer, tol):
scene = gs.Scene(
vis_options=gs.options.VisOptions(
# rendered_envs_idx=(0, 1, 2),
Expand Down Expand Up @@ -247,6 +247,7 @@ def test_batched_offscreen_rendering(show_viewer, tol):
)
scene.build(n_envs=3, env_spacing=(2.0, 2.0))

cam.start_recording()
for _ in range(7):
dofs_lower_bound, dofs_upper_bound = robot.get_dofs_limit()
qpos = dofs_lower_bound + (dofs_upper_bound - dofs_lower_bound) * torch.rand(robot.n_qs)
Expand All @@ -270,6 +271,7 @@ def test_batched_offscreen_rendering(show_viewer, tol):

for i in range(3):
assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol)
cam.stop_recording(save_to_filename=(tmp_path / "video.mp4"))


@pytest.mark.required
Expand Down