Source code for sapien.utils.viewer.viewer

import os
from pathlib import Path

import numpy as np
import sapien
from sapien import Entity, Scene
from sapien import internal_renderer as R
from sapien.render import (
    RenderSystem,
    RenderWindow,
    SapienRenderer,
    get_viewer_shader_dir,
)

from .articulation_window import ArticulationWindow
from .contact_window import ContactWindow
from .control_window import ControlWindow
from .entity_window import EntityWindow
from .imgui_ini import imgui_ini
from .path_window import PathWindow
from .plugin import Plugin
from .render_window import RenderOptionsWindow
from .scene_window import SceneWindow
from .setting_window import SettingWindow
from .transform_window import TransformWindow

# from .keyframe_window import KeyframeWindow


[docs] class Viewer: def __init__( self, renderer: SapienRenderer = None, shader_dir="", resolutions=(1920, 1080), plugins=[ PathWindow(), ContactWindow(), SettingWindow(), TransformWindow(), RenderOptionsWindow(), ControlWindow(), SceneWindow(), EntityWindow(), ArticulationWindow(), # KeyframeWindow(), ], ): filename = sapien.render.get_imgui_ini_filename() if not filename: filename = "imgui.ini" if not os.path.exists(filename): Path(filename).parent.mkdir(parents=True, exist_ok=True) with open(filename, "w") as f: f.write(imgui_ini) if renderer is None: renderer = SapienRenderer() self.renderer_context = renderer._internal_context if not shader_dir: self.shader_dir = get_viewer_shader_dir() else: self.shader_dir = shader_dir resolution = np.array(resolutions).flatten()[:2] self.scenes = [] self.window = RenderWindow(*resolution, self.shader_dir) self.window.set_focus_callback(self.focus_change) self.window.set_drop_callback(self.drop) self.paused = False self.render_target = "Color" self.plugins = plugins self.init_plugins(plugins) self._selected_entity_visibility = 0.5 @property def render_scene(self): return self.window._internal_scene @property def scene(self) -> sapien.Scene: if len(self.scenes) == 1: return self.scenes[0] return None @property def cameras(self): if self.scene: return self.scene.render_system.cameras return []
[docs] def drop(self, files): if not self.scene: return builder = self.scene.create_actor_builder() for f in files: builder.add_visual_from_file(f) builder.build_kinematic("dropped file")
[docs] def focus_change(self, focused): for plugin in self.plugins: plugin.notify_window_focus_change(focused)
[docs] def init_plugins(self, plugins): for plugin in plugins: assert isinstance(plugin, Plugin) plugin.init(self)
[docs] def set_scenes(self, scenes, offsets=None): if offsets is None: side = int(np.ceil(len(scenes) ** 0.5)) idx = np.arange(len(scenes)) offsets = np.stack([idx // side, idx % side, np.zeros_like(idx)], axis=1) if self.scenes: camera_pose = self.window.get_camera_pose() self.clear_scene() else: camera_pose = sapien.Pose([-2, 0, 0.5]) self.selected_entity = None self.scenes = scenes if len(scenes) == 0: self.window.set_scene(None) elif len(scenes) == 1: self.window.set_scene(scenes[0]) self.scene_offset = {scenes[0]: np.array([0, 0, 0])} else: self.window.set_scenes(scenes, offsets) self.scene_offset = dict((s, o) for s, o in zip(scenes, offsets)) self.window.set_camera_parameters(0.1, 1000, np.pi / 2) self.set_camera_pose(camera_pose) for plugin in self.plugins: plugin.notify_scene_change()
[docs] def get_entity_viewer_pose(self, entity): return sapien.Pose(self.scene_offset[entity.scene]) * entity.pose
[docs] def set_scene(self, scene: Scene): self.set_scenes([scene])
# if self.scene is not None: # camera_pose = self.window.get_camera_pose() # self.clear_scene() # else: # camera_pose = sapien.Pose([-2, 0, 0.5]) # self.selected_entity = None # self.scene = scene # # self.system = self.scene.render_system # self.window.set_scene(scene) # self.window.set_camera_parameters(0.1, 1000, np.pi / 2) # self.set_camera_pose(camera_pose) # for plugin in self.plugins: # plugin.notify_scene_change()
[docs] def clear_scene(self): for plugin in self.plugins: plugin.clear_scene()
@property def closed(self): return self.window is None
[docs] def close(self): for plugin in self.plugins: plugin.close() self.selected_entity = None self.scenes = [] # self.system = None self.window = None self.plugins = [] self.renderer_context = None
[docs] def set_camera_pose(self, pose): self.window.set_camera_pose(pose) self.notify_render_update()
[docs] def notify_render_update(self): """notify the viewer that the camera is moved""" self.render_updated = True
[docs] def reset_notifications(self): self.render_updated = False
[docs] def render(self): if self.window.should_close: self.close() if self.closed: return while True: if self.window.should_close: break if not self.paused or self.render_updated: self.window.update_render() # self.system.step() self.reset_notifications() for plugin in self.plugins: plugin.before_render() ui_windows = [] for plugin in self.plugins: ui_windows += plugin.get_ui_windows() self.window.render(self.render_target, ui_windows) for plugin in self.plugins: plugin.after_render() if not self.paused: break
[docs] def select_entity(self, entity: Entity): if self.selected_entity == entity: return # reset previous selected entity if self.selected_entity is not None: for c in self.selected_entity.components: if isinstance(c, sapien.render.RenderBodyComponent): c.visibility = 1 self.selected_entity = entity # update selected entity if self.selected_entity is not None: for c in self.selected_entity.components: if isinstance(c, sapien.render.RenderBodyComponent): c.visibility = self.selected_entity_visibility for plugin in self.plugins: plugin.notify_selected_entity_change() self.render_scene.force_rebuild()
@property def selected_entity_visibility(self): return self._selected_entity_visibility @selected_entity_visibility.setter def selected_entity_visibility(self, v): self.notify_render_update() self._selected_entity_visibility = v if self.selected_entity is not None: for c in self.selected_entity.components: if isinstance(c, sapien.render.RenderBodyComponent): c.visibility = self.selected_entity_visibility @property def resolution(self): return self.window.size @resolution.setter def resolution(self, res): self.window.resize(res[0], res[1])
[docs] def add_bounding_box(self, pose, half_size, color): vertices = np.array( [ [1, -1, -1], [1, 1, -1], [-1, 1, -1], [-1, -1, -1], [1, -1, 1], [1, 1, 1], [-1, 1, 1], [-1, -1, 1], ] ) lines = [0, 1, 1, 2, 2, 3, 3, 0, 4, 5, 5, 6, 6, 7, 7, 4, 0, 4, 1, 5, 2, 6, 3, 7] vertices = vertices[lines] colors = np.ones((vertices.shape[0], 4)) * [*color[:3], 1] lineset = self.renderer_context.create_line_set(vertices, colors) # render_scene: R.Scene = self.system._internal_scene box = self.render_scene.add_line_set(lineset) box.set_position(pose.p) box.set_rotation(pose.q) box.set_scale(half_size) return box
[docs] def update_bounding_box(self, box, pose, half_size): box.set_position(pose.p) box.set_rotation(pose.q) box.set_scale(half_size)
[docs] def remove_bounding_box(self, box): # render_scene: R.Scene = self.system._internal_scene self.render_scene.remove_node(box)
[docs] def draw_aabb(self, lower, upper, color): pose = sapien.Pose((lower + upper) / 2) half_size = (upper - lower) / 2 return self.add_bounding_box(pose, half_size, color)
[docs] def update_aabb(self, aabb, lower, upper): pose = sapien.Pose((lower + upper) / 2) half_size = (upper - lower) / 2 self.update_bounding_box(aabb, pose, half_size)
@property def control_window(self) -> ControlWindow: for plugin in self.plugins: if isinstance(plugin, ControlWindow): return plugin return None
[docs] def loop(self, physx_steps=0): """ A convenience method for opening a temporary viewer for a scene. Simply call scene.create_viewer().loop() """ while not self.closed: for _ in range(physx_steps): self.scene.physx_system.step() self.scene.update_render() self.render()
[docs] def register_click_handler(self, handler): ...
[docs] def set_camera_xyz(self, x, y, z): ...
[docs] def set_camera_rpy(self, r, p, y): ...
[docs] def focus_entity(self, entity): ...
[docs] def focus_camera(self, camera): ...