Source code for sapien.utils.viewer.transform_window
import numpy as np
import sapien
from sapien import internal_renderer as R
from .plugin import Plugin
[docs]
class TransformWindow(Plugin):
def __init__(self):
self.reset()
[docs]
def reset(self):
self.ui_window = None
self.gizmo = None
self._gizmo_pose = sapien.Pose()
self.ghost_objects = []
self.ik_articulation = None
self.ik_enabled = True
self.ik_errpr = None
self.ik_result = None
self.ik_success = False
self.follow = False
self.display_ghosts = False
self.enabled = False
self.ui_move_group = None
self.move_group_joints = []
self.move_group_selection = []
self.pinocchio_model = None
@property
def gizmo_matrix(self):
return self._gizmo_pose.to_transformation_matrix()
@gizmo_matrix.setter
def gizmo_matrix(self, v):
self._gizmo_pose = sapien.Pose(v)
self.update_ghost_objects()
self.follow = False
[docs]
def get_articulation(self, entity):
if not entity:
return None
for c in entity.components:
if isinstance(c, sapien.physx.PhysxArticulationLinkComponent):
return c.articulation
return None
[docs]
def notify_selected_entity_change(self):
if not self.scene:
return
self.clear_ghost_objects()
if self.viewer.selected_entity is None:
self._gizmo_pose = sapien.Pose()
self.follow = False
else:
self._gizmo_pose = self.selected_entity.pose
self.follow = True
art = self.get_articulation(self.selected_entity)
if art is not None:
if self.ik_articulation == art:
return
self.ik_articulation = art
self.pinocchio_model: sapien.PinocchioModel = (
self.ik_articulation.create_pinocchio_model()
)
self.move_group_joints = [
j.name
for j in self.ik_articulation.get_joints()
if j.get_dof() != 0
]
self.move_group_selection = [True] * len(self.move_group_joints)
self.ui_move_group.remove_children()
for i, n in enumerate(self.move_group_joints):
self.ui_move_group.append(
R.UICheckbox().Label(n).Bind(self.move_group_selection, i)
)
[docs]
def compute_ik(self):
if (
self.selected_entity is not None
and self.ik_enabled
and self.get_articulation(self.selected_entity)
):
link_idx = self.ik_articulation.get_links().index(
next(
c
for c in self.selected_entity.components
if isinstance(c, sapien.physx.PhysxArticulationLinkComponent)
)
)
mask = np.array(
[
self.move_group_selection[j]
for j in range(len(self.move_group_joints))
]
).astype(int)
pose = self.ik_articulation.pose.inv() * self._gizmo_pose
result, success, error = self.pinocchio_model.compute_inverse_kinematics(
link_idx,
pose,
initial_qpos=self.ik_articulation.get_qpos(),
active_qmask=mask,
max_iterations=100,
)
return result, success, error
@property
def scene(self):
return self.viewer.scene
@property
def selected_entity(self):
return self.viewer.selected_entity
[docs]
def clear_ghost_objects(self):
render_scene: R.Scene = self.viewer.render_scene
for obj in self.ghost_objects:
render_scene.remove_node(obj)
self.ghost_objects = []
self.viewer.notify_render_update()
[docs]
def refresh_ghost_objects(self):
render_scene: R.Scene = self.viewer.render_scene
self.clear_ghost_objects()
if self.selected_entity is None:
return
render_body = None
articulation_link = None
for c in self.selected_entity.components:
if isinstance(c, sapien.render.RenderBodyComponent):
render_body = c
if isinstance(c, sapien.physx.PhysxArticulationLinkComponent):
articulation_link = c
if render_body is None:
return
elif articulation_link is not None:
link: sapien.LinkBase = self.selected_entity
link2world = link.pose
articulation = articulation_link.articulation
for l in articulation.get_links():
new_node = render_scene.add_node()
for body in [
c
for c in l.entity.components
if isinstance(c, sapien.render.RenderBodyComponent)
]:
render_node = body._internal_node
for obj in render_node.children:
new_obj = render_scene.add_object(obj.model, new_node)
new_obj.set_position(obj.position)
new_obj.set_rotation(obj.rotation)
new_obj.set_scale(obj.scale)
new_obj.transparency = 0.7
new_obj.set_segmentation(obj.get_segmentation())
new_node.set_position(l.pose.p)
new_node.set_rotation(l.pose.q)
self.ghost_objects.append(new_node)
else:
entity2world = self.selected_entity.pose
render_node = render_body._internal_node
new_node = render_scene.add_node()
for obj in render_node.children:
new_obj = render_scene.add_object(obj.model, new_node)
new_obj.set_position(obj.position)
new_obj.set_rotation(obj.rotation)
new_obj.set_scale(obj.scale)
new_obj.transparency = 0.7
new_node.set_position(self.selected_entity.pose.p)
new_node.set_rotation(self.selected_entity.pose.q)
self.ghost_objects.append(new_node)
self.viewer.notify_render_update()
[docs]
def update_ghost_objects(self):
if self.selected_entity is None:
return
if not self.ghost_objects:
self.refresh_ghost_objects()
art = self.get_articulation(self.selected_entity)
if art is not None:
if self.ik_enabled:
# IK
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.ghost_objects):
pose = art.pose * self.pinocchio_model.get_link_pose(idx)
obj.set_position(pose.p)
obj.set_rotation(pose.q)
else:
link = self.selected_entity
link2world = link.pose
for l, node in zip(art.get_links(), self.ghost_objects):
newlink2world = self._gizmo_pose
l2world = l.pose
l2link = link2world.inv() * l2world
newl2world = newlink2world * l2link
node.set_position(newl2world.p)
node.set_rotation(newl2world.q)
return
else:
for node in self.ghost_objects:
node.set_position(self._gizmo_pose.p)
node.set_rotation(self._gizmo_pose.q)
self.viewer.notify_render_update()
[docs]
def teleport(self, _):
try:
art = self.get_articulation(self.selected_entity)
if art:
link: sapien.LinkBase = self.selected_entity
if self.ik_enabled:
art.set_qpos(self.ik_result)
else:
link2world = link.pose
newlink2world = self._gizmo_pose
l2world = art.pose
l2link = link2world.inv() * l2world
newl2world = newlink2world * l2link
art.set_root_pose(newl2world)
else:
self.selected_entity.set_pose(self._gizmo_pose)
self.viewer.notify_render_update()
except AttributeError:
pass
[docs]
def build(self):
if self.scene is None:
self.ui_window = None
return
if self.follow and self.selected_entity and self.ghost_objects:
self._gizmo_pose = self.selected_entity.pose
self.update_ghost_objects()
if not self.ui_window:
self.ui_move_group = R.UISection().Label("Move Group")
self.gizmo = R.UIGizmo().Bind(self, "gizmo_matrix")
self.ui_window = (
R.UIWindow()
.Label("Transform")
.Pos(10, 10)
.Size(400, 400)
.append(
R.UICheckbox().Label("Enabled").Bind(self, "enabled"),
R.UIConditional()
.Bind(self, "enabled")
.append(
self.gizmo,
R.UIConditional()
.Bind(lambda: self.selected_entity is not None)
.append(
R.UIConditional()
.Bind(
lambda: self.get_articulation(self.selected_entity)
is not None
)
.append(
R.UICheckbox().Label("IK").Bind(self, "ik_enabled"),
self.ui_move_group,
),
R.UIButton().Label("Teleport").Callback(self.teleport),
),
),
)
)
proj = self.viewer.window.get_camera_projection_matrix()
view = (
(
self.viewer.window.get_camera_pose()
* sapien.Pose([0, 0, 0], [-0.5, -0.5, 0.5, 0.5])
)
.inv()
.to_transformation_matrix()
)
self.gizmo.CameraMatrices(view, proj)
if self.selected_entity is not None:
self.gizmo.Matrix(self.selected_entity.pose.to_transformation_matrix())
else:
self.gizmo.Matrix(np.eye(4))