from ..core import renderer as R
from ..core import (
ActorBase,
Pose,
VulkanRenderer,
Scene,
VulkanWindow,
ArticulationBase,
Joint,
LinkBase,
LightEntity,
PointLightEntity,
DirectionalLightEntity,
SpotLightEntity,
CameraEntity
)
from transforms3d.quaternions import axangle2quat as aa
from transforms3d.euler import quat2euler
from transforms3d.quaternions import qmult, mat2quat, rotate_vector, qinverse
import numpy as np
import os
imgui_ini = """
[Window][DockSpace Demo]
Pos=0,0
Size=1024,768
Collapsed=0
[Window][Actor/Entity]
Pos=726,23
Size=298,335
Collapsed=0
DockId=0x00000007,0
[Window][Control]
Pos=0,23
Size=248,368
Collapsed=0
DockId=0x00000003,0
[Window][Scene Hierarchy]
Pos=0,393
Size=248,314
Collapsed=0
DockId=0x00000004,0
[Window][Articulation]
Pos=726,360
Size=298,347
Collapsed=0
DockId=0x00000008,0
[Window][Info]
Pos=0,709
Size=1024,59
Collapsed=0
DockId=0x0000000A,0
[Window][IK]
Pos=726,360
Size=298,347
Collapsed=0
DockId=0x00000008,1
[Window][Debug##Default]
Pos=60,60
Size=400,400
Collapsed=0
[Docking][Data]
DockSpace ID=0x4BBE4C7A Window=0x4647B76E Pos=0,23 Size=1024,745 Split=Y
DockNode ID=0x00000009 Parent=0x4BBE4C7A SizeRef=1024,684 Split=X
DockNode ID=0x00000005 Parent=0x00000009 SizeRef=724,747 Split=X
DockNode ID=0x00000001 Parent=0x00000005 SizeRef=248,747 Split=Y Selected=0x9A68760C
DockNode ID=0x00000003 Parent=0x00000001 SizeRef=399,368 Selected=0x226615D7
DockNode ID=0x00000004 Parent=0x00000001 SizeRef=399,314 Selected=0x9A68760C
DockNode ID=0x00000002 Parent=0x00000005 SizeRef=474,747 CentralNode=1
DockNode ID=0x00000006 Parent=0x00000009 SizeRef=298,747 Split=Y Selected=0x85B479FD
DockNode ID=0x00000007 Parent=0x00000006 SizeRef=121,335 Selected=0x0A9B993B
DockNode ID=0x00000008 Parent=0x00000006 SizeRef=121,347 Selected=0x816C7EAB
DockNode ID=0x0000000A Parent=0x4BBE4C7A SizeRef=1024,59 Selected=0x6BBB9E69
"""
[docs]class FPSCameraController:
def __init__(self, window: VulkanWindow):
self.window = window
self.forward = np.array([1, 0, 0])
self.up = np.array([0, 0, 1])
self.left = np.cross(self.up, self.forward)
self.initial_rotation = mat2quat(
np.array([-self.left, self.up, -self.forward]).T
)
self.xyz = np.zeros(3)
self.rpy = np.zeros(3)
[docs] def setRPY(self, roll, pitch, yaw):
self.rpy = np.array([roll, pitch, yaw])
self.update()
[docs] def setXYZ(self, x, y, z):
self.xyz = np.array([x, y, z])
self.update()
[docs] def move(self, forward, left, up):
q = qmult(
qmult(aa(self.up, -self.rpy[2]), aa(self.left, -self.rpy[1])),
aa(self.forward, self.rpy[0]),
)
self.xyz = self.xyz + (
rotate_vector(self.forward, q) * forward
+ rotate_vector(self.left, q) * left
+ rotate_vector(self.up, q) * up
)
self.update()
[docs] def rotate(self, roll, pitch, yaw):
self.rpy = self.rpy + np.array([roll, pitch, yaw])
self.update()
[docs] def update(self):
self.rpy[1] = np.clip(self.rpy[1], -1.57, 1.57)
if self.rpy[2] >= 3.15:
self.rpy[2] = self.rpy[2] - 2 * np.pi
elif self.rpy[2] <= -3.15:
self.rpy[2] = self.rpy[2] + 2 * np.pi
rot = qmult(
qmult(
qmult(aa(self.up, -self.rpy[2]), aa(self.left, -self.rpy[1])),
aa(self.forward, self.rpy[0]),
),
self.initial_rotation,
)
self.window.set_camera_rotation(rot)
self.window.set_camera_position(self.xyz)
[docs]class ArcRotateCameraController:
def __init__(self, window: VulkanWindow):
self.window = window
self.forward = np.array([1, 0, 0])
self.up = np.array([0, 0, 1])
self.left = np.cross(self.up, self.forward)
self.initial_rotation = mat2quat(
np.array([-self.left, self.up, -self.forward]).T
)
self.center = np.zeros(3)
self.yaw = 0
self.pitch = 0
self.radius = 1
[docs] def set_center(self, center):
self.center = np.array(center)
self.update()
[docs] def rotate_yaw_pitch(self, yaw, pitch):
self.yaw += yaw
self.pitch += pitch
self.update()
[docs] def set_yaw_pitch(self, yaw, pitch):
self.yaw = yaw
self.pitch = pitch
self.update()
[docs] def zoom(self, zoom_in):
self.radius -= zoom_in
self.radius = max(0.1, self.radius)
self.radius = min(100, self.radius)
self.update()
[docs] def set_zoom(self, zoom):
self.radius = zoom
self.update()
[docs] def update(self):
rot = qmult(
qmult(aa(self.up, self.yaw), aa(self.left, self.pitch)),
self.initial_rotation,
)
pos = self.center - self.radius * rotate_vector(np.array([0, 0, -1]), rot)
self.window.set_camera_rotation(rot)
self.window.set_camera_position(pos)
[docs]class Viewer(object):
def __init__(
self,
renderer: VulkanRenderer,
shader_dir="",
resolutions=((1024, 768), (800, 600), (1920, 1080)),
):
if not os.path.exists("imgui.ini"):
with open("imgui.ini", "w") as f:
f.write(imgui_ini)
self.shader_dir = shader_dir
self.renderer = renderer
self.renderer_context: R.Context = renderer._internal_context
self.create_visual_models()
self.scene = None
self.window = None
self.resolution = None
self.resolutions = None
self.set_window_resolutions(resolutions)
self.axes = None
self.show_axes = True
self.axes_scale = 0.1
self.selected_entity = None
self.focused_entity = None
self.paused = False
self.target_name = "Color"
self.single_step = False
self.focused_camera = None
self.cameras = None
self.camera_ui = None
self.move_speed = 0.05
self.rotate_speed = 0.005
self.scroll_speed = 0.5
self.selection_opacity = 0.7
self.scene_window = None
self.control_window = None
self.info_window = None
self.ik_window = None
self.gizmo = None
self.ik_enabled = False
self.ik_display_objects = []
self.move_group_joints = []
self.move_group_selection = {}
self.key_stack = ""
self.initialize_key_action_map()
self.mode = "normal"
self.display_object = None
self.coordinate_axes_mode = "Origin"
self.immediate_mode = False
self.camera_linesets = []
self._show_camera_linesets = True
def _clear_camera_linesets(self):
if self.scene is None:
return
rs = self.scene.renderer_scene
for n in self.camera_linesets:
rs._internal_scene.remove_node(n)
self.camera_linesets = []
def _update_camera_linesets(self):
if self.scene is None:
return
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
cameras = self.scene.get_cameras()
if len(self.camera_linesets) != len(cameras):
self._clear_camera_linesets()
for c in self.cameras:
self.camera_linesets.append(
render_scene.add_line_set(self.camera_lineset)
)
for lineset, camera in zip(self.camera_linesets, cameras):
lineset: R.LineSetObject
mat = camera.get_model_matrix()
lineset.set_position(mat[:3, 3])
lineset.set_rotation(mat2quat(mat[:3, :3]))
scaley = np.tan(camera.fovy / 2)
scalex = np.tan(camera.fovx / 2)
lineset.set_scale(np.array([scalex, scaley, 1]) * 0.3)
[docs] def create_visual_models(self):
self.cone = self.renderer_context.create_cone_mesh(16)
self.capsule = self.renderer_context.create_capsule_mesh(0.1, 0.5, 16, 4)
self.mat_red = self.renderer_context.create_material(
[0, 0, 0, 1], [1, 0, 0, 1], 0, 0, 0
)
self.mat_green = self.renderer_context.create_material(
[0, 0, 0, 1], [0, 1, 0, 1], 0, 0, 0
)
self.mat_blue = self.renderer_context.create_material(
[0, 0, 0, 1], [0, 0, 1, 1], 0, 0, 0
)
self.mat_cyan = self.renderer_context.create_material(
[0, 0, 0, 1], [0, 1, 1, 1], 0, 0, 0
)
self.mat_magenta = self.renderer_context.create_material(
[0, 0, 0, 1], [1, 0, 1, 1], 0, 0, 0
)
self.red_cone = self.renderer_context.create_model([self.cone], [self.mat_red])
self.green_cone = self.renderer_context.create_model(
[self.cone], [self.mat_green]
)
self.blue_cone = self.renderer_context.create_model(
[self.cone], [self.mat_blue]
)
self.red_capsule = self.renderer_context.create_model(
[self.capsule], [self.mat_red]
)
self.green_capsule = self.renderer_context.create_model(
[self.capsule], [self.mat_green]
)
self.blue_capsule = self.renderer_context.create_model(
[self.capsule], [self.mat_blue]
)
self.cyan_capsule = self.renderer_context.create_model(
[self.capsule], [self.mat_cyan]
)
self.magenta_capsule = self.renderer_context.create_model(
[self.capsule], [self.mat_magenta]
)
self.camera_lineset = self.renderer_context.create_line_set(
[
0,
0,
0,
1,
1,
-1,
0,
0,
0,
-1,
1,
-1,
0,
0,
0,
1,
-1,
-1,
0,
0,
0,
-1,
-1,
-1,
1,
1,
-1,
1,
-1,
-1,
1,
-1,
-1,
-1,
-1,
-1,
-1,
-1,
-1,
-1,
1,
-1,
-1,
1,
-1,
1,
1,
-1,
1,
1.2,
-1,
0,
2,
-1,
0,
2,
-1,
-1,
1.2,
-1,
-1,
1.2,
-1,
1,
1.2,
-1,
],
[0.9254901960784314, 0.5764705882352941, 0.18823529411764706, 1] * 22,
)
def _create_coordinate_axes(self):
assert self.scene is not None
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
node = render_scene.add_node()
obj = render_scene.add_object(self.red_cone, node)
obj.set_scale([0.5, 0.2, 0.2])
obj.set_position([1, 0, 0])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
obj = render_scene.add_object(self.red_capsule, node)
obj.set_position([0.5, 0, 0])
obj.set_scale([1.02, 1.02, 1.02])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
obj = render_scene.add_object(self.green_cone, node)
obj.set_scale([0.5, 0.2, 0.2])
obj.set_position([0, 1, 0])
obj.set_rotation([0.7071068, 0, 0, 0.7071068])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
obj = render_scene.add_object(self.green_capsule, node)
obj.set_position([0, 0.5, 0])
obj.set_rotation([0.7071068, 0, 0, 0.7071068])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
obj = render_scene.add_object(self.blue_cone, node)
obj.set_scale([0.5, 0.2, 0.2])
obj.set_position([0, 0, 1])
obj.set_rotation([0, 0.7071068, 0, 0.7071068])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
obj = render_scene.add_object(self.blue_capsule, node)
obj.set_position([0, 0, 0.5])
obj.set_rotation([0, 0.7071068, 0, 0.7071068])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
return node
def _create_grab_axes(self):
assert self.scene is not None
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
grab_axes = [
render_scene.add_object(model)
for model in [self.red_capsule, self.green_capsule, self.blue_capsule]
]
for obj in grab_axes:
obj.set_position([0, 0, 0])
obj.set_scale([100, 0.1, 0.1])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
return grab_axes
def _create_joint_axes(self):
assert self.scene is not None
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
joint_axes = [
render_scene.add_object(self.magenta_capsule),
render_scene.add_object(self.cyan_capsule),
]
for obj in joint_axes:
obj.set_position([0, 0, 0])
obj.set_scale([5, 0.1, 0.1])
obj.shading_mode = 2
obj.cast_shadow = False
obj.transparency = 1
return joint_axes
[docs] def create_visual_objects(self):
if hasattr(self, "coordinate_axes") and self.coordinate_axes:
self.scene.renderer_scene._internal_scene.remove_node(self.coordinate_axes)
for n in self.grab_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
for n in self.joint_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
self.coordinate_axes = self._create_coordinate_axes()
self.grab_axes = self._create_grab_axes()
self.joint_axes = self._create_joint_axes()
[docs] def enter_mode(self, name):
if self.mode == name:
return
self.leave_mode(self.mode)
self.mode = name
if name == "grab":
self.window.cursor = False
else:
self.window.cursor = True
[docs] def leave_mode(self, name):
if name == "grab":
for obj in self.grab_axes:
obj.transparency = 1
if self.display_object:
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
render_scene.remove_node(self.display_object)
self.display_object = None
elif name == "rotate":
for obj in self.grab_axes:
obj.transparency = 1
if self.display_object:
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
render_scene.remove_node(self.display_object)
self.display_object = None
[docs] def add_display_object(self):
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
if not isinstance(self.selected_entity, ActorBase):
self.display_object = render_scene.add_node()
self.display_object.set_position(self.selected_entity.pose.p)
self.display_object.set_rotation(self.selected_entity.pose.q)
return
if self.display_object:
render_scene.remove_node(self.display_object)
self.display_object = None
self.display_object = render_scene.add_node()
selected2world = self.selected_entity.pose
for body in self.selected_entity.get_visual_bodies():
for obj in body._internal_objects:
scale = obj.scale
obj2world = Pose(obj.position, obj.rotation)
obj2selected = selected2world.inv() * obj2world
new_obj = render_scene.add_object(obj.model, self.display_object)
new_obj.set_position(obj2selected.p)
new_obj.set_rotation(obj2selected.q)
new_obj.set_scale(scale)
new_obj.transparency = 0.1
self.display_object.set_position(selected2world.p)
self.display_object.set_rotation(selected2world.q)
[docs] def initialize_key_action_map(self):
x2y = np.array([0.7071068, 0, 0, 0.7071068])
x2z = np.array([0.7071068, 0, 0.7071068, 0])
def f():
if self.selected_entity:
self.focus_entity(self.selected_entity)
def r():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.add_display_object()
point = self.world_space_to_screen_space(self.selected_entity.pose.p)
axis = self.screen_space_to_world_space(
[point[0], point[1], 1]
) - self.screen_space_to_world_space([point[0], point[1], 0])
axis = axis / np.linalg.norm(axis)
self.rotate_axis = np.array(axis)
self.rotate_direction = 1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def rx():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation([1, 0, 0, 0])
self.add_display_object()
self.rotate_axis = np.array([1, 0, 0])
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def rxx():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation(self.selected_entity.pose.q)
self.add_display_object()
self.rotate_axis = np.array(
rotate_vector([1, 0, 0], self.selected_entity.pose.q)
)
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def rxxx():
self.key_stack = "r"
r()
def ry():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(x2y)
self.add_display_object()
self.rotate_axis = np.array([0, 1, 0])
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def ryy():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(qmult(self.selected_entity.pose.q, x2y))
self.add_display_object()
self.rotate_axis = np.array(
rotate_vector([0, 1, 0], self.selected_entity.pose.q)
)
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def ryyy():
self.key_stack = "r"
r()
def rz():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(x2z)
self.add_display_object()
self.rotate_axis = np.array([0, 0, 1])
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def rzz():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("rotate")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(qmult(self.selected_entity.pose.q, x2z))
self.add_display_object()
self.rotate_axis = np.array(
rotate_vector([0, 0, 1], self.selected_entity.pose.q)
)
screen_vector = self.world_space_to_screen_space(
self.selected_entity.pose.p
) - self.world_space_to_screen_space(
self.selected_entity.pose.p + self.rotate_axis
)
if screen_vector[2] <= 0:
self.rotate_direction = 1
else:
self.rotate_direction = -1
self.rotate_initial_mouse_position = np.array(self.window.mouse_position)
self.rotate_initial_rotation = np.array(self.selected_entity.pose.q)
self.rotate_screen_center = self.world_space_to_screen_space(
self.selected_entity.pose.p
)[:2]
def rzzz():
self.key_stack = "r"
r()
def g():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.add_display_object()
self.grab_axis = None
self.grab_plane = None
def gx():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation([1, 0, 0, 0])
self.add_display_object()
self.grab_axis = np.array([1, 0, 0])
self.grab_plane = None
def gxx():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation(self.selected_entity.pose.q)
self.add_display_object()
self.grab_axis = rotate_vector([1, 0, 0], self.selected_entity.pose.q)
self.grab_plane = None
def gxxx():
self.key_stack = "g"
g()
def gX():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 0
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(x2y)
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(x2z)
self.add_display_object()
self.grab_axis = None
self.grab_plane = np.array([1, 0, 0])
def gXX():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 0
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(qmult(self.selected_entity.pose.q, x2y))
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(qmult(self.selected_entity.pose.q, x2z))
self.add_display_object()
self.grab_axis = None
self.grab_plane = rotate_vector(
np.array([1, 0, 0]), self.selected_entity.pose.q
)
def gXXX():
self.key_stack = "g"
g()
def gy():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(x2y)
self.add_display_object()
self.grab_axis = np.array([0, 1, 0])
self.grab_plane = None
def gyy():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(qmult(self.selected_entity.pose.q, x2y))
self.add_display_object()
self.grab_axis = rotate_vector([0, 1, 0], self.selected_entity.pose.q)
self.grab_plane = None
def gyyy():
self.key_stack = "g"
g()
def gY():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation([1, 0, 0, 0])
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(x2z)
self.add_display_object()
self.grab_axis = None
self.grab_plane = np.array([0, 1, 0])
def gYY():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation(
qmult(self.selected_entity.pose.q, [1, 0, 0, 0])
)
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(qmult(self.selected_entity.pose.q, x2z))
self.add_display_object()
self.grab_axis = None
self.grab_plane = rotate_vector(
np.array([0, 1, 0]), self.selected_entity.pose.q
)
def gYYY():
self.key_stack = "g"
g()
def gz():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(x2z)
self.add_display_object()
self.grab_axis = np.array([0, 0, 1])
self.grab_plane = None
def gzz():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 1
self.grab_axes[1].transparency = 1
self.grab_axes[2].transparency = 0
self.grab_axes[2].set_position(self.selected_entity.pose.p)
self.grab_axes[2].set_rotation(qmult(self.selected_entity.pose.q, x2z))
self.add_display_object()
self.grab_axis = rotate_vector([0, 0, 1], self.selected_entity.pose.q)
self.grab_plane = None
def gzzz():
self.key_stack = "g"
g()
def gZ():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation([1, 0, 0, 0])
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(x2y)
self.add_display_object()
self.grab_axis = None
self.grab_plane = np.array([0, 0, 1])
def gZZ():
if not self.selected_entity:
self.key_stack = ""
return
self.enter_mode("grab")
self.grab_axes[0].transparency = 0
self.grab_axes[1].transparency = 0
self.grab_axes[2].transparency = 1
self.grab_axes[0].set_position(self.selected_entity.pose.p)
self.grab_axes[0].set_rotation(
qmult(self.selected_entity.pose.q, [1, 0, 0, 0])
)
self.grab_axes[1].set_position(self.selected_entity.pose.p)
self.grab_axes[1].set_rotation(qmult(self.selected_entity.pose.q, x2y))
self.add_display_object()
self.grab_axis = None
self.grab_plane = rotate_vector(
np.array([0, 0, 1]), self.selected_entity.pose.q
)
def gZZZ():
self.key_stack = "g"
g()
self.key_press_action_map = {
"f": f,
"r": r,
"rx": rx,
"rxx": rxx,
"rxxx": rxxx,
"ry": ry,
"ryy": ryy,
"ryyy": ryyy,
"rz": rz,
"rzz": rzz,
"rzzz": rzzz,
"g": g,
"gx": gx,
"gxx": gxx,
"gxxx": gxxx,
"gX": gX,
"gXX": gXX,
"gXXX": gXXX,
"gy": gy,
"gyy": gyy,
"gyyy": gyyy,
"gY": gY,
"gYY": gYY,
"gYYY": gYYY,
"gz": gz,
"gzz": gzz,
"gzzz": gzzz,
"gZ": gZ,
"gZZ": gZZ,
"gZZZ": gZZZ,
}
def w():
self.enter_mode("normal")
speed_mod = 0.1 if self.window.shift else 1
self.focus_entity(None)
self.focus_camera(None)
self.fps_camera_controller.move(self.move_speed * speed_mod, 0, 0)
self.key_stack = ""
def s():
self.enter_mode("normal")
speed_mod = 0.1 if self.window.shift else 1
self.focus_entity(None)
self.focus_camera(None)
self.fps_camera_controller.move(-self.move_speed * speed_mod, 0, 0)
self.key_stack = ""
def a():
self.enter_mode("normal")
speed_mod = 0.1 if self.window.shift else 1
self.focus_entity(None)
self.focus_camera(None)
self.fps_camera_controller.move(0, self.move_speed * speed_mod, 0)
self.key_stack = ""
def d():
self.enter_mode("normal")
speed_mod = 0.1 if self.window.shift else 1
self.focus_entity(None)
self.focus_camera(None)
self.fps_camera_controller.move(0, -self.move_speed * speed_mod, 0)
self.key_stack = ""
self.key_down_action_map = {"w": w, "s": s, "a": a, "d": d}
[docs] def key_press_action(self, key):
if len(key) == 1 and ord("a") <= ord(key) <= ord("z") and self.window.shift:
key = key.upper()
if self.key_stack + key in self.key_press_action_map:
self.key_stack += key
self.key_press_action_map[self.key_stack]()
return
if self.key_stack == "":
if key in self.key_press_action_map:
self.key_stack = key
self.key_press_action_map[key]()
return
self.key_stack = self.key_stack[:-1]
self.key_press_action(key)
[docs] def key_down_action(self, key):
if key in self.key_down_action_map:
self.key_down_action_map[key]()
[docs] def set_fovy(self, fovy):
self.fovy = fovy
self.window.set_camera_parameters(0.1, 100, fovy)
[docs] def set_window_resolutions(self, resolutions):
assert len(resolutions)
for r in resolutions:
assert len(r) == 2
self.window = self.renderer.create_window(
resolutions[0][0], resolutions[0][1], self.shader_dir
)
self.resolution = resolutions[0]
self.resolutions = resolutions
[docs] def build_control_window(self):
if not self.control_window:
self.cameras = self.scene.get_cameras()
self.camera_ui = (
R.UIOptions()
.Style("select")
.Label("Name##camera_name")
.Index(
0
if self.focused_camera is None
else self.cameras.index(self.focused_camera) + 1
)
.Items(
["None"]
+ [x.get_name() + f"##{i}" for i, x in enumerate(self.cameras)]
)
.Callback(
lambda p: self.focus_camera(
self.cameras[p.index - 1] if p.index > 0 else None
)
)
)
self.control_window = (
R.UIWindow()
.Label("Control")
.Pos(10, 10)
.Size(400, 400)
.append(
R.UISameLine().append(
R.UICheckbox()
.Label("Pause")
.Callback(lambda p: self.toggle_pause(p.checked)),
R.UIButton()
.Label("Single Step")
.Callback(lambda p: self.step_button()),
),
R.UIDisplayText().Text("Camera Speed"),
R.UISliderFloat()
.Min(0.01)
.Max(1)
.Value(self.move_speed)
.Label("Move")
.Callback(lambda w: self.set_move_speed(w.value)),
R.UISliderFloat()
.Min(0.001)
.Max(0.01)
.Value(self.rotate_speed)
.Label("Rotate")
.Callback(lambda w: self.set_rotate_speed(w.value)),
R.UISliderFloat()
.Min(0.1)
.Max(1)
.Value(self.scroll_speed)
.Label("Scroll")
.Callback(lambda w: self.set_scroll_speed(w.value)),
R.UIDisplayText().Text("Camera"),
self.camera_ui,
R.UIDisplayText().Text("Display Settings"),
R.UISliderAngle()
.Min(1)
.Max(179)
.Value(self.fovy)
.Label("Fov Y")
.Callback(lambda w: self.set_fovy(w.value)),
R.UIOptions()
.Style("select")
.Label("Render Target")
.Index(0)
.Items(
["Color"]
+ [x for x in self.window.display_target_names if x != "Color"]
)
.Callback(lambda p: self.set_target(p.value)),
R.UIOptions()
.Style("select")
.Label("Resolution")
.Index(0)
.Items(["{}x{}".format(r[0], r[1]) for r in self.resolutions])
.Callback(lambda p: self.set_resolution(p.index)),
R.UIDisplayText().Text("Actor Selection"),
R.UICheckbox()
.Label("Coordinate Axes")
.Checked(True)
.Callback(lambda p: self.toggle_axes(p.checked)),
R.UICheckbox()
.Label("Camera Display")
.Checked(True)
.Callback(lambda p: self.toggle_camera_lines(p.checked)),
R.UIOptions()
.Style("select")
.Label("Axes Mode")
.Index(0)
.Items(["Origin", "Center of Mass"])
.Callback(lambda p: self.set_coordinate_axes_mode(p.value)),
R.UISliderFloat()
.Label("Axes Scale")
.Min(0)
.Max(1)
.Value(self.axes_scale)
.Callback(lambda p: self.update_coordinate_axes_scale(p.value)),
R.UISliderFloat()
.Label("Opacity")
.Min(0)
.Max(1)
.Value(self.selection_opacity)
.Callback(lambda p: self.set_selection_opacity(p.value)),
R.UICheckbox()
.Label("Immediate Move")
.Checked(self.immediate_mode)
.Callback(lambda p: self.set_immediate_move(p.checked)),
R.UIButton()
.Label("Take Screenshot")
.Callback(self.take_screenshot),
R.UIDisplayText().Text("FPS: {:.2f}".format(self.window.fps)),
)
)
self.control_window.get_children()[-1].Text(
"FPS: {:.2f}".format(self.window.fps)
)
[docs] def enable_ik(self, enable):
self.ik_enabled = enable
self.refresh_ik()
self.refresh_ik_display_objects()
[docs] def build_ik_window(self):
if not self.ik_window:
self.gizmo = R.UIGizmo().Matrix(np.eye(4))
self.move_group = R.UISection().Label("Move Group")
self.ik_window = (
R.UIWindow()
.Label("IK")
.Pos(10, 10)
.Size(400, 400)
.append(
R.UISameLine().append(
R.UICheckbox()
.Label("Enable IK")
.Callback(lambda c: self.enable_ik(c.checked)),
R.UIButton().Label("Go!").Callback(self.execute_ik),
),
self.move_group,
)
)
if self.ik_enabled:
if self.ik_window.get_children()[-1] != self.gizmo:
self.ik_window.append(self.gizmo)
elif self.ik_window.get_children()[-1] == self.gizmo:
self.ik_window = (
R.UIWindow()
.Label("IK")
.Pos(10, 10)
.Size(400, 400)
.append(
R.UISameLine().append(
R.UICheckbox()
.Label("Enable IK")
.Callback(lambda c: self.enable_ik(c.checked)),
R.UIButton().Label("Go!").Callback(self.execute_ik),
),
self.move_group,
)
)
[docs] def take_screenshot(self, _):
picture = self.window.get_float_texture("Color")
for i in range(100000000):
n = f"sapien_screenshot_{i}.png"
if os.path.exists(n):
continue
from PIL import Image
Image.fromarray((picture.clip(0, 1) * 255).astype(np.uint8)).save(n)
break
[docs] def set_selection_opacity(self, opacity):
self.selection_opacity = opacity
[docs] def set_resolution(self, index):
self.resolution = self.resolutions[index]
self.window.resize(*self.resolution)
[docs] def build_scene_window(self):
assert self.scene
if not self.scene_window:
self.scene_window = (
R.UIWindow()
.Pos(410, 10)
.Size(400, 400)
.Label("Scene Hierarchy")
.append(
R.UITreeNode()
.Label("World")
.append(
R.UITreeNode().Label("Actors"),
R.UITreeNode().Label("Articulations"),
R.UITreeNode().Label("Lights"),
),
)
)
atree, arttree, ltree = self.scene_window.get_children()[0].get_children()
atree: R.UITreeNode
arttree: R.UITreeNode
ltree: R.UITreeNode
atree.remove_children()
arttree.remove_children()
ltree.remove_children()
for i, actor in enumerate(self.scene.get_all_actors()):
atree.append(
R.UISelectable()
.Label(
"{}##actor{}".format(actor.name if actor.name else "(no name)", i)
)
.Selected(self.selected_entity == actor)
.Callback((lambda link: lambda _: self.select_entity(link))(actor))
)
for i, art in enumerate(self.scene.get_all_articulations()):
art_node = R.UITreeNode().Label(
"{}##art{}".format(art.name if art.name else "(no name)", i)
)
for j, link in enumerate(art.get_links()):
art_node.append(
R.UISelectable()
.Label(
"{}##link{}".format(link.name if link.name else "(no name)", j)
)
.Selected(self.selected_entity == link)
.Callback((lambda link: lambda _: self.select_entity(link))(link))
)
arttree.append(art_node)
for i, light in enumerate(self.scene.get_all_lights()):
ltree.append(
R.UISelectable()
.Label(
"{}##light{}".format(light.name if light.name else "(no name)", i)
)
.Selected(self.selected_entity == light)
.Callback((lambda light: lambda _: self.select_entity(light))(light))
)
[docs] def build_actor_window(self):
self.actor_window = R.UIWindow().Label("Actor/Entity")
if not self.selected_entity:
self.actor_window.append(
R.UIDisplayText().Text("No actor/entity selected.")
)
return
if isinstance(self.selected_entity, ActorBase):
actor = self.selected_entity
self.actor_window.append(
R.UIDisplayText().Text("Name: {}".format(actor.name)),
R.UIDisplayText().Text("Class: {}".format(actor.classname)),
R.UIDisplayText().Text("Id: {}".format(actor.id)),
)
self.actor_window.append(
R.UIDisplayText().Text("Position"),
R.UIInputFloat()
.Label("x##actorpx")
.Value(actor.pose.p[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorpy")
.Value(actor.pose.p[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorpz")
.Value(actor.pose.p[2])
.ReadOnly(True),
R.UIDisplayText().Text("Rotation"),
R.UIInputFloat()
.Label("w##actorqw")
.Value(actor.pose.q[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("x##actorqx")
.Value(actor.pose.q[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorqy")
.Value(actor.pose.q[2])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorqz")
.Value(actor.pose.q[3])
.ReadOnly(True),
R.UISameLine().append(
R.UIButton()
.Label("Show")
.Callback(
(lambda actor: lambda _: actor.render_collision(True))(actor)
),
R.UIButton()
.Label("Hide")
.Callback(
(lambda actor: lambda _: actor.render_collision(False))(actor)
),
R.UIDisplayText().Text("Collision"),
),
)
if actor.classname in ["Actor", "Link"]:
self.actor_window.append(
R.UIInputFloat().Label("Mass").Value(actor.mass).ReadOnly(True)
)
# collision shapes
shape_section = R.UISection().Label("Collision Shapes")
self.actor_window.append(shape_section)
shapes = actor.get_collision_shapes()
for shape_idx, shape in enumerate(shapes):
c0, c1, c2, c3 = shape.get_collision_groups()
shape_pose = shape.get_local_pose()
mat = shape.get_physical_material()
shape_info = (
R.UITreeNode()
.Label("{}##{}".format(shape.type, shape_idx))
.append(
R.UIDisplayText().Text(
"Contact offset: {:.3g}".format(shape.contact_offset)
),
R.UIDisplayText().Text(
"Rest offset: {:.3g}".format(shape.rest_offset)
),
R.UIDisplayText().Text(
"Patch radius: {:.3g}".format(shape.patch_radius)
),
R.UIDisplayText().Text(
"Min path radius: {:.3g}".format(shape.min_patch_radius)
),
R.UICheckbox().Label("Is trigger").Checked(shape.is_trigger),
R.UIDisplayText().Text(
"Static friction: {:.3g}".format(mat.get_static_friction())
),
R.UIDisplayText().Text(
"Dynamic friction: {:.3g}".format(
mat.get_dynamic_friction()
)
),
R.UIDisplayText().Text(
"Restitution: {:.3g}".format(mat.get_restitution())
),
R.UIDisplayText().Text("Collision groups:"),
R.UIDisplayText().Text(" 0x{:08x} 0x{:08x}".format(c0, c1)),
R.UIDisplayText().Text(" 0x{:08x} 0x{:08x}".format(c2, c3)),
R.UIDisplayText().Text("Local position"),
R.UIInputFloat()
.Label("x##actorpx")
.Value(shape_pose.p[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorpy")
.Value(shape_pose.p[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorpz")
.Value(shape_pose.p[2])
.ReadOnly(True),
R.UIDisplayText().Text("Local rotation"),
R.UIInputFloat()
.Label("w##actorqw")
.Value(shape_pose.q[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("x##actorqx")
.Value(shape_pose.q[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorqy")
.Value(shape_pose.q[2])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorqz")
.Value(shape_pose.q[3])
.ReadOnly(True),
)
)
shape_section.append(shape_info)
if shape.type == "sphere":
shape_info.append(
R.UIDisplayText().Text(
"Sphere radius: {:.3g}".format(shape.geometry.radius)
)
)
elif shape.type == "capsule":
shape_info.append(
R.UIDisplayText().Text(
"Capsule radius: {:.3g}".format(shape.geometry.radius)
),
R.UIDisplayText().Text(
"Capsule half length: {:.3g}".format(
shape.geometry.half_length
)
),
)
elif shape.type == "box":
x, y, z = shape.geometry.half_lengths
shape_info.append(
R.UIDisplayText().Text(
"Box half lengths: {:.3g} {:.3g} {:.3g}".format(x, y, z)
)
)
elif shape.type == "convex_mesh":
x, y, z = shape.geometry.scale
shape_info.append(
R.UIDisplayText().Text(
"Mesh scale: {:.3g} {:.3g} {:.3g}".format(x, y, z)
)
)
elif shape.type == "nonconvex_mesh":
x, y, z = shape.geometry.scale
shape_info.append(
R.UIDisplayText().Text(
"Mesh scale: {:.3g} {:.3g} {:.3g}".format(x, y, z)
)
)
# render shapes
body_section = R.UISection().Label("Visual Bodies")
self.actor_window.append(body_section)
bodies = actor.get_visual_bodies()
for body_idx, body in enumerate(bodies):
body_info = R.UITreeNode().Label("{}##{}".format(body.type, body_idx))
body_section.append(body_info)
if body.type == "sphere":
body_info.append(
R.UIDisplayText().Text("Radius: {:.3g}".format(body.radius))
)
elif body.type == "capsule":
body_info.append(
R.UIDisplayText().Text("Radius: {:.3g}".format(body.radius)),
R.UIDisplayText().Text(
"Half length: {:.3g}".format(body.half_length)
),
)
elif body.type == "box":
x, y, z = body.half_lengths
body_info.append(
R.UIDisplayText().Text(
"Half extents: {:.3g} {:.3g} {:.3g}".format(x, y, z)
)
)
elif body.type == "mesh":
x, y, z = body.scale
body_info.append(
R.UIDisplayText().Text(
"Scale: {:.3g} {:.3g} {:.3g}".format(x, y, z)
)
)
body_info.append(
R.UIDisplayText().Text("Visual id: {}".format(body.visual_id))
)
shapes = body.get_render_shapes()
for shape_idx, shape in enumerate(shapes):
mat = shape.material
dtex = (
(mat.diffuse_texture_filename or "(has texture)")
if mat.diffuse_texture
else "(no texture)"
)
rtex = (
(mat.roughness_texture_filename or "(has texture)")
if mat.roughness_texture
else "(no texture)"
)
mtex = (
(mat.metallic_texture_filename or "(has texture)")
if mat.metallic_texture
else "(no texture)"
)
ntex = (
(mat.normal_texture_filename or "(has texture)")
if mat.normal_texture
else "(no texture)"
)
etex = (
(mat.emission_texture_filename or "(has texture)")
if mat.emission_texture
else "(no texture)"
)
body_info.append(
R.UITreeNode()
.Label("material {}##{}".format(shape_idx, body_idx))
.append(
R.UIInputFloat4()
.Label("Diffuse")
.ReadOnly(True)
.Value(mat.base_color),
R.UIInputText()
.ReadOnly(True)
.Value(dtex)
.Label("##dtex{}_{}".format(shape_idx, body_idx)),
R.UIInputFloat4()
.Label("Emission")
.ReadOnly(True)
.Value(mat.emission),
R.UIInputText()
.ReadOnly(True)
.Value(etex)
.Label("##etex{}_{}".format(shape_idx, body_idx)),
R.UIInputFloat()
.Label("Roughness")
.ReadOnly(True)
.Value(mat.roughness),
R.UIInputText()
.ReadOnly(True)
.Value(rtex)
.Label("##rtex{}_{}".format(shape_idx, body_idx)),
R.UIInputFloat()
.Label("Metallic")
.ReadOnly(True)
.Value(mat.metallic),
R.UIInputText()
.ReadOnly(True)
.Value(mtex)
.Label("##mtex{}_{}".format(shape_idx, body_idx)),
R.UIInputFloat()
.Label("Specular")
.ReadOnly(True)
.Value(mat.specular),
R.UIDisplayText().Text("Normal map"),
R.UIInputText()
.ReadOnly(True)
.Value(ntex)
.Label("##ntex{}_{}".format(shape_idx, body_idx)),
)
)
if isinstance(self.selected_entity, LightEntity):
light = self.selected_entity
self.actor_window.append(
R.UIDisplayText().Text("Name: {}".format(light.name)),
R.UIDisplayText().Text("Class: {}".format(light.classname)),
)
def set_shadow(light, enable):
light.shadow = enable
self.actor_window.append(
R.UICheckbox()
.Label("Shadow")
.Checked(light.shadow)
.Callback((lambda light: lambda p: set_shadow(light, p.checked))(light))
)
self.actor_window.append(
R.UIDisplayText().Text("Position"),
R.UIInputFloat()
.Label("x##actorpx")
.Value(light.pose.p[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorpy")
.Value(light.pose.p[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorpz")
.Value(light.pose.p[2])
.ReadOnly(True),
R.UIDisplayText().Text("Rotation"),
R.UIInputFloat()
.Label("w##actorqw")
.Value(light.pose.q[0])
.ReadOnly(True),
R.UIInputFloat()
.Label("x##actorqx")
.Value(light.pose.q[1])
.ReadOnly(True),
R.UIInputFloat()
.Label("y##actorqy")
.Value(light.pose.q[2])
.ReadOnly(True),
R.UIInputFloat()
.Label("z##actorqz")
.Value(light.pose.q[3])
.ReadOnly(True),
)
self.actor_window.append(
R.UIInputFloat().Label("Near").Value(light.shadow_near).ReadOnly(True),
R.UIInputFloat().Label("Far").Value(light.shadow_far).ReadOnly(True),
)
if light.classname == "PointLightEntity":
light: PointLightEntity
pass
elif light.classname == "DirectionalLightEntity":
self.actor_window.append(
R.UIInputFloat()
.Label("Half Size")
.Value(light.shadow_half_size)
.ReadOnly(True),
)
pass
elif light.classname == "SpotLightEntity":
pass
[docs] def build_articulation_window(self):
self.articulation_window = R.UIWindow().Label("Articulation")
if (
not self.selected_entity
or not isinstance(self.selected_entity, ActorBase)
or self.selected_entity.classname not in ["Link", "KinematicLink"]
):
self.articulation_window.append(
R.UIDisplayText().Text("No articulation selected.")
)
return
art = self.selected_entity.get_articulation()
art: ArticulationBase
self.articulation_window.append(
R.UIDisplayText().Text(
"Name: {}".format(art.name if art.name else "(no name)")
),
R.UIDisplayText().Text("Class: {}".format(art.classname)),
R.UIDisplayText().Text("Base Link Id: {}".format(art.get_links()[0].id)),
)
uijoints = R.UISection().Label("Joints")
joints = []
for j in art.get_joints():
if j.get_dof() > 0:
joints.append(j)
def wrapper(art, i, qpos):
def callback(slider):
qpos[i] = slider.value
art.set_qpos(qpos)
return callback
qpos = art.get_qpos()
for i, (q, j) in enumerate(zip(qpos, joints)):
line = R.UISameLine()
line.append(
R.UISliderFloat()
.Label(j.name + "##joint_{}".format(i))
.Min(max(j.get_limits()[0][0], -20))
.Max(min(j.get_limits()[0][1], 20))
.Value(q)
.Callback(wrapper(art, i, qpos)),
)
if art.classname == "Articulation":
j: Joint
line.append(
R.UITreeNode()
.Label("##joint_expand_{}".format(i))
.append(
R.UISliderFloat()
.Label("Drive Target##{}".format(i))
.Min(max(j.get_limits()[0][0], -20))
.Max(min(j.get_limits()[0][1], 20))
.Value(j.get_drive_target())
.Callback((lambda j: lambda p: j.set_drive_target(p.value))(j)),
R.UIInputFloat()
.Label("Damping##{}".format(i))
.Value(j.damping)
.Callback(
(
lambda j: lambda p: j.set_drive_property(
j.stiffness, p.value, j.force_limit, j.drive_mode,
)
)(j)
),
R.UIInputFloat()
.Label("Stiffness##{}".format(i))
.Value(j.stiffness)
.Callback(
(
lambda j: lambda p: j.set_drive_property(
p.value, j.damping, j.force_limit, j.drive_mode,
)
)(j)
),
R.UIInputFloat()
.Label("Force Limit##{}".format(i))
.Value(j.force_limit)
.Callback(
(
lambda j: lambda p: j.set_drive_property(
j.stiffness, j.damping, p.value, j.drive_mode,
)
)(j)
),
R.UIInputFloat()
.Label("Friction##{}".format(i))
.Value(j.friction)
.Callback((lambda j: lambda p: j.set_friction(p.value))(j)),
R.UICheckbox()
.Label("Acceleration##{}".format(i))
.Checked(j.drive_mode == "acceleration")
.Callback(
(
lambda j: lambda p: j.set_drive_property(
j.stiffness,
j.damping,
j.force_limit,
"acceleration" if p.checked else "force",
)
)(j)
),
)
)
uijoints.append(line)
self.articulation_window.append(uijoints)
def wrapper(art, show):
def show_link_collision(_):
for link in art.get_links():
link.render_collision(show)
return show_link_collision
self.articulation_window.append(
R.UISameLine().append(
R.UIButton().Label("Show").Callback(wrapper(art, True)),
R.UIButton().Label("Hide").Callback(wrapper(art, False)),
R.UIDisplayText().Text("Collision"),
),
)
[docs] def build_info_window(self):
if not self.info_window:
self.info_window = R.UIWindow().Label("Info").append(R.UIDisplayText())
self.info_window.get_children()[0].Text("-".join(self.key_stack))
[docs] def set_move_speed(self, x):
self.move_speed = x
[docs] def set_rotate_speed(self, x):
self.rotate_speed = x
[docs] def set_scene(self, scene: Scene):
if hasattr(self, "scene") and self.scene:
if hasattr(self, "coordinate_axes") and self.coordinate_axes:
self.scene.renderer_scene._internal_scene.remove_node(
self.coordinate_axes
)
for n in self.grab_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
for n in self.joint_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
self.coordinate_axes = None
self.grab_axes = None
self.joint_axes = None
self._clear_camera_linesets()
self.scene = scene
self.window.set_scene(scene)
self.fps_camera_controller = FPSCameraController(self.window)
self.arc_camera_controller = ArcRotateCameraController(self.window)
self.create_visual_objects()
self.toggle_axes(True)
self.set_fovy(np.pi / 2)
[docs] def set_camera_xyz(self, x, y, z):
self.fps_camera_controller.setXYZ(x, y, z)
[docs] def set_camera_rpy(self, r, p, y):
self.fps_camera_controller.setRPY(r, p, y)
[docs] def toggle_pause(self, paused):
self.paused = paused
[docs] def toggle_axes(self, show):
self.show_axes = show
for c in self.coordinate_axes.children:
if show:
c.transparency = 0
else:
c.transparency = 1
[docs] def toggle_camera_lines(self, show):
self._show_camera_linesets = show
if not show:
self._clear_camera_linesets()
[docs] def set_coordinate_axes_mode(self, mode):
self.coordinate_axes_mode = mode
[docs] def set_target(self, name):
self.target_name = name
[docs] def find_actor(self, id):
actors = self.scene.get_all_actors()
for actor in actors:
if actor.id == id:
return actor
for a in self.scene.get_all_articulations():
for link in a.get_links():
if link.id == id:
return link
@property
def closed(self):
return self.window is None
[docs] def close(self):
self.gizmo = None
self.move_group = None
if hasattr(self, "coordinate_axes") and self.coordinate_axes:
self.scene.renderer_scene._internal_scene.remove_node(self.coordinate_axes)
for n in self.grab_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
for n in self.joint_axes:
self.scene.renderer_scene._internal_scene.remove_node(n)
self.coordinate_axes = None
self.grab_axes = None
self.joint_axes = None
self._clear_camera_linesets()
self.axes = None
self.scene = None
self.fps_camera_controller = None
self.arc_camera_controller = None
self.window = None
self.camera_ui = None
self.control_window = None
self.scene_window = None
self.actor_window = None
self.articulation_window = None
self.info_window = None
self.ik_window = None
[docs] def focus_entity(self, actor):
if actor == self.focused_entity:
return
self.focused_entity = actor
if self.focused_entity is not None:
self.focus_camera(None)
pos = self.window.get_camera_position()
rot = self.window.get_camera_rotation()
x, y, z = rotate_vector([0, 0, -1], rot)
yaw = -np.arctan2(y, x)
pitch = np.arctan2(z, np.linalg.norm([x, y]))
self.arc_camera_controller.set_yaw_pitch(-yaw, -pitch)
self.arc_camera_controller.set_center(actor.pose.p)
self.arc_camera_controller.set_zoom(np.linalg.norm(actor.pose.p - pos))
else:
rot = self.window.get_camera_rotation()
x, y, z = rotate_vector([0, 0, -1], rot)
yaw = -np.arctan2(y, x)
pitch = np.arctan2(z, np.linalg.norm([x, y]))
self.fps_camera_controller.setXYZ(*self.window.get_camera_position())
self.fps_camera_controller.setRPY(0, pitch, yaw)
[docs] def select_entity(self, entity):
self.enter_mode("normal")
self.key_stack = ""
if entity:
if self.selected_entity and isinstance(self.selected_entity, ActorBase):
for v in self.selected_entity.get_visual_bodies():
v.set_visibility(1)
self.selected_entity = entity
if isinstance(self.selected_entity, ActorBase):
for v in entity.get_visual_bodies():
v.set_visibility(self.selection_opacity)
else:
if self.selected_entity and isinstance(self.selected_entity, ActorBase):
for v in self.selected_entity.get_visual_bodies():
v.set_visibility(1)
self.selected_entity = None
self.update_coordinate_axes()
self.refresh_ik()
self.refresh_ik_display_objects()
[docs] def refresh_ik(self):
if (
self.selected_entity
and isinstance(self.selected_entity, LinkBase)
and self.ik_enabled
):
self.pinocchio_model = (
self.selected_entity.get_articulation().create_pinocchio_model()
)
self.gizmo.Matrix(self.selected_entity.pose.to_transformation_matrix())
self.move_group_joints = [
j.name
for j in self.selected_entity.get_articulation().get_joints()
if j.get_dof() != 0
]
for n in self.move_group_joints:
if n not in self.move_group_selection:
self.move_group_selection[n] = True
self.move_group.remove_children()
def select_move_group(name, select):
self.move_group_selection[name] = select
for n in self.move_group_joints:
self.move_group.append(
R.UICheckbox()
.Label(n)
.Checked(self.move_group_selection[n])
.Callback((lambda n: lambda c: select_move_group(n, c.checked))(n))
)
else:
self.pinocchio_model = None
self.move_group.remove_children()
[docs] def clear_ik_display_objects(self):
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
for obj in self.ik_display_objects:
render_scene.remove_node(obj)
self.ik_display_objects = []
[docs] def refresh_ik_display_objects(self):
rs = self.scene.renderer_scene
render_scene: R.Scene = rs._internal_scene
self.clear_ik_display_objects()
if (
self.selected_entity is not None
and self.ik_enabled
and isinstance(self.selected_entity, LinkBase)
):
for link in self.selected_entity.get_articulation().get_links():
link2world = link.pose
display_obj = render_scene.add_node()
for body in link.get_visual_bodies():
for obj in body._internal_objects:
scale = obj.scale
obj2world = Pose(obj.position, obj.rotation)
obj2selected = link2world.inv() * obj2world
new_obj = render_scene.add_object(obj.model, display_obj)
new_obj.set_position(obj2selected.p)
new_obj.set_rotation(obj2selected.q)
new_obj.set_scale(scale)
new_obj.transparency = 0.1
display_obj.set_position(link.pose.p)
display_obj.set_rotation(link.pose.q)
self.ik_display_objects.append(display_obj)
[docs] def update_ik_display_objects(self):
if (
self.selected_entity is not None
and self.ik_enabled
and isinstance(self.selected_entity, LinkBase)
):
link_idx = (
self.selected_entity.get_articulation()
.get_links()
.index(self.selected_entity)
)
pose = Pose.from_transformation_matrix(self.gizmo.matrix)
self.ik_display_objects[link_idx].set_position(pose.p)
self.ik_display_objects[link_idx].set_rotation(pose.q)
result, success, error = self.compute_ik()
self.ik_result = result
self.ik_success = success
self.ik_errpr = error
self.pinocchio_model.compute_forward_kinematics(result)
for idx, obj in enumerate(self.ik_display_objects):
pose = self.selected_entity.get_articulation().pose * self.pinocchio_model.get_link_pose(
idx
)
obj.set_position(pose.p)
obj.set_rotation(pose.q)
[docs] def compute_ik(self):
if (
self.selected_entity is not None
and self.ik_enabled
and isinstance(self.selected_entity, LinkBase)
):
a = self.selected_entity.get_articulation()
link_idx = a.get_links().index(self.selected_entity)
mask = np.array(
[self.move_group_selection[j] for j in self.move_group_joints]
).astype(int)
pose = a.pose.inv() * Pose.from_transformation_matrix(self.gizmo.matrix)
result, success, error = self.pinocchio_model.compute_inverse_kinematics(
link_idx,
pose,
initial_qpos=a.get_qpos(),
active_qmask=mask,
max_iterations=100,
)
return result, success, error
[docs] def execute_ik(self, _):
if (
self.selected_entity is not None
and self.ik_enabled
and isinstance(self.selected_entity, LinkBase)
):
if hasattr(self, "ik_result"):
mask = np.array(
[self.move_group_selection[j] for j in self.move_group_joints]
).astype(int)
a = self.selected_entity.get_articulation()
target = a.get_drive_target()
target = self.ik_result * mask + target * (1 - mask)
# self.selected_entity.get_articulation().set_qpos(target)
self.selected_entity.get_articulation().set_drive_target(target)
[docs] @staticmethod
def get_camera_pose(camera: CameraEntity):
"""Get the camera pose in the Sapien world."""
opengl_pose = camera.get_model_matrix() # opengl camera-> sapien world
# sapien camera -> opengl camera
sapien2opengl = np.array(
[
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[-1, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
]
)
cam_pose = Pose.from_transformation_matrix(opengl_pose @ sapien2opengl)
return cam_pose
[docs] def focus_camera(self, camera: CameraEntity):
if self.focused_camera == camera:
return
self.focused_camera = camera
if self.focused_camera is not None:
self.focus_entity(None)
if self.camera_ui is not None:
# Lazy check if any camera has changed
assert self.cameras == self.scene.get_cameras(), "Cameras have changed"
index = (self.cameras.index(camera) + 1) if camera is not None else 0
self.camera_ui.Index(index)
[docs] def update_coordinate_axes_scale(self, scale):
self.axes_scale = scale
self.update_coordinate_axes()
[docs] def update_coordinate_axes(self):
if self.selected_entity:
self.coordinate_axes.set_scale([self.axes_scale] * 3)
if self.coordinate_axes_mode == "Origin":
self.coordinate_axes.set_position(self.selected_entity.pose.p)
self.coordinate_axes.set_rotation(self.selected_entity.pose.q)
elif self.coordinate_axes_mode == "Center of Mass":
if isinstance(
self.selected_entity, ActorBase
) and self.selected_entity.classname in [
"Actor",
"Link",
"KinematicLink",
]:
pose = (
self.selected_entity.pose
* self.selected_entity.cmass_local_pose
)
else:
pose = self.selected_entity.pose
self.coordinate_axes.set_position(pose.p)
self.coordinate_axes.set_rotation(pose.q)
else:
self.coordinate_axes.set_position([0, 0, 0])
self.coordinate_axes.set_rotation([1, 0, 0, 0])
self.coordinate_axes.set_scale([self.axes_scale] * 3)
[docs] def update_joint_axis(self):
if (
self.selected_entity
and isinstance(self.selected_entity, ActorBase)
and "Link" in self.selected_entity.classname
):
link: LinkBase = self.selected_entity
j = link.get_articulation().get_joints()[link.get_index()]
if j.type not in ["revolute", "prismatic"]:
for x in self.joint_axes:
x.transparency = 1
j2c = j.get_pose_in_child()
c2w = link.get_pose()
j2w = c2w * j2c
if j.type == "prismatic":
j2w.set_p(c2w.p)
self.joint_axes[1].set_position(j2w.p)
self.joint_axes[1].set_rotation(j2w.q)
self.joint_axes[1].transparency = 0 if self.show_axes else 1
self.joint_axes[0].transparency = 1
elif j.type == "revolute":
self.joint_axes[0].set_position(j2w.p)
self.joint_axes[0].set_rotation(j2w.q)
self.joint_axes[0].transparency = 0 if self.show_axes else 1
self.joint_axes[1].transparency = 1
else:
self.joint_axes[0].transparency = 1
self.joint_axes[1].transparency = 1
[docs] def render(self):
if self.closed:
return
self.set_fovy(self.fovy)
while True:
self.scene.update_render()
self.build_control_window()
self.build_scene_window()
self.build_actor_window()
self.build_articulation_window()
self.build_info_window()
self.build_ik_window()
proj = self.window.get_camera_projection_matrix()
view = (
Pose(
self.window.get_camera_position(), self.window.get_camera_rotation()
)
.inv()
.to_transformation_matrix()
)
self.gizmo.CameraMatrices(view, proj)
self.window.render(
self.target_name,
[
self.control_window,
self.scene_window,
self.actor_window,
self.articulation_window,
self.info_window,
self.ik_window,
],
)
if self.mode == "grab":
if self.window.mouse_click(0) or self.immediate_mode:
new_pose = Pose(
self.display_object.position, self.display_object.rotation
)
if isinstance(self.selected_entity, LightEntity):
self.selected_entity.set_pose(new_pose)
elif self.selected_entity.classname == "Actor":
self.selected_entity.set_pose(new_pose)
elif self.selected_entity.classname in ["Link", "KinematicLink"]:
old_pose = self.selected_entity.pose
a = self.selected_entity.get_articulation()
a.set_root_pose(new_pose * old_pose.inv() * a.get_root_pose())
if self.window.mouse_click(0):
self.enter_mode("normal")
elif self.mode == "rotate":
if self.window.mouse_click(0) or self.immediate_mode:
new_pose = Pose(
self.display_object.position, self.display_object.rotation
)
if isinstance(self.selected_entity, LightEntity):
self.selected_entity.set_pose(new_pose)
if self.selected_entity.classname == "Actor":
self.selected_entity.set_pose(new_pose)
elif self.selected_entity.classname in ["Link", "KinematicLink"]:
old_pose = self.selected_entity.pose
a = self.selected_entity.get_articulation()
a.set_root_pose(new_pose * old_pose.inv() * a.get_root_pose())
if self.window.mouse_click(0):
self.enter_mode("normal")
elif self.mode == "normal":
if self.window.mouse_click(0):
mx, my = self.window.mouse_position
if not self.is_mouse_available(mx, my):
print("[W] Mouse not available")
continue
ww, wh = self.window.size
tw, th = self.window.get_target_size("Segmentation")
mx = mx * tw / ww
my = my * th / wh
pixel = self.window.get_uint32_texture_pixel(
"Segmentation", int(mx), int(my)
)
actor = self.find_actor(pixel[1])
self.select_entity(actor)
self.update_coordinate_axes()
self.update_joint_axis()
self.update_ik_display_objects()
if self._show_camera_linesets:
self._update_camera_linesets()
speed_mod = 1
if self.window.shift:
speed_mod *= 0.1
if self.window.key_press("esc"):
self.enter_mode("normal")
for key in "frgxyz":
if self.window.key_press(key):
self.key_press_action(key)
for key in "wasd":
if self.window.key_down(key):
self.key_down_action(key)
if self.mode == "grab":
x, y = self.window.mouse_delta
vec = self.camera_space_to_world_space(np.array([x, -y, 0]))
if self.grab_axis is not None:
vec = vec @ self.grab_axis * self.grab_axis
elif self.grab_plane is not None:
v = self.world_space_to_camera_space(self.grab_plane)
vec = self.camera_space_to_world_space(
[x, -y, (v[0] * x + v[1] * (-y)) / (-v[2])]
)
self.display_object.set_position(
self.display_object.position + vec * 0.01
)
self.coordinate_axes.set_position(self.display_object.position)
self.coordinate_axes.set_rotation(self.display_object.rotation)
if self.mode == "rotate":
mouse_position = np.array(self.window.mouse_position)
dir1 = self.rotate_initial_mouse_position - self.rotate_screen_center
dir2 = mouse_position - self.rotate_screen_center
angle = (
self.rotate_direction
* np.sign(np.cross(dir1, dir2))
* np.arccos((dir1 @ dir2) / ((dir1 @ dir1) * (dir2 @ dir2)) ** 0.5)
)
self.display_object.set_rotation(
qmult(aa(self.rotate_axis, angle), self.rotate_initial_rotation)
)
self.coordinate_axes.set_position(self.display_object.position)
self.coordinate_axes.set_rotation(self.display_object.rotation)
if self.mode == "normal":
if self.window.mouse_down(1):
x, y = self.window.mouse_delta
if self.focused_entity:
self.arc_camera_controller.rotate_yaw_pitch(
-self.rotate_speed * speed_mod * x,
self.rotate_speed * speed_mod * y,
)
else:
self.fps_camera_controller.rotate(
0,
-self.rotate_speed * speed_mod * y,
self.rotate_speed * speed_mod * x,
)
if self.window.mouse_down(2):
x, y = self.window.mouse_delta
self.focus_entity(None)
self.fps_camera_controller.move(
0,
self.rotate_speed * speed_mod * x,
self.rotate_speed * speed_mod * y,
)
wx, wy = self.window.mouse_wheel_delta
if wx != 0:
if self.focused_entity:
self.arc_camera_controller.zoom(
self.scroll_speed * speed_mod * wx
)
else:
self.fps_camera_controller.move(
self.scroll_speed * speed_mod * wx, 0, 0
)
if self.focused_entity:
self.arc_camera_controller.set_center(self.focused_entity.pose.p)
elif self.focused_camera:
cam_pose = self.get_camera_pose(self.focused_camera)
rpy = quat2euler(cam_pose.q)
self.set_camera_xyz(*cam_pose.p)
self.set_camera_rpy(rpy[0], -rpy[1], -rpy[2])
if self.window.key_down("q") or self.window.should_close:
self.close()
return
if not self.paused or (self.paused and self.single_step):
self.single_step = False
break
[docs] def is_mouse_available(self, mx, my):
w, h = self.window.size
# print(f"[I] windowSize: {w, h}; mousePose: {mx, my}")
return mx >= 0 and my >= 0
[docs] def camera_space_to_world_space(self, vec):
return rotate_vector(vec, self.window.get_camera_rotation())
[docs] def world_space_to_camera_space(self, vec):
return rotate_vector(vec, qinverse(self.window.get_camera_rotation()))
[docs] def world_space_to_screen_space(self, point):
q = qinverse(self.window.get_camera_rotation())
point = rotate_vector(point, q) - rotate_vector(
self.window.get_camera_position(), q
)
proj = self.window.get_camera_projection_matrix()
point = proj @ np.concatenate((point, [1]))
point = point[:3] / point[3]
point[:2] = (point[:2] * 0.5 + 0.5) * np.array(self.window.size)
return point
[docs] def screen_space_to_world_space(self, point):
proj = self.window.get_camera_projection_matrix()
point[:2] = (point[:2] / np.array(self.window.size)) * 2 - 1
point = np.linalg.inv(proj) @ np.concatenate((point, [1]))
point = point[:3] / point[3] # screen space
q = self.window.get_camera_rotation()
point = rotate_vector(point, q) + self.window.get_camera_position()
return point