import os
import platform
import numpy as np
import sapien
from sapien import internal_renderer as R
from .plugin import Plugin
[docs]
class Keyframe(R.UIKeyframe):
def __init__(self, serialized_scene, frame=0):
super().__init__()
self.serialized_scene = serialized_scene
self._frame = frame
[docs]
def frame(self):
return self._frame
[docs]
def set_frame(self, frame):
self._frame = frame
def __repr__(self):
return "Keyframe(scene, frame={})".format(self._frame)
[docs]
class Duration(R.UIDuration):
DEFAULT_DEFINITION = """import sapien
import numpy as np
class Reward:
def __init__(self, env, scene: sapien.Scene):
self.env = env
self.scene = scene
def compute(self):
return 0
"""
def __init__(self, keyframe0, keyframe1, name="", definition=DEFAULT_DEFINITION):
super().__init__()
self._keyframe0 = keyframe0
self._keyframe1 = keyframe1
self._name = name
self.definition = definition
[docs]
def keyframe0(self):
return self._keyframe0
[docs]
def keyframe1(self):
return self._keyframe1
[docs]
def set_name(self, name):
self._name = name
[docs]
def name(self):
return self._name
[docs]
def serialize_keyframes(keyframes, durations):
s_keyframes = [(f.serialized_scene, f._frame) for f in keyframes]
f2i = dict((f, i) for i, f in enumerate(keyframes))
s_durations = [
(f2i[d._keyframe0], f2i[d._keyframe1], d._name, d.definition) for d in durations
]
return s_keyframes, s_durations
[docs]
def deserialize_keyframes(state):
s_keyframes, s_durations = state
keyframes = [Keyframe(*f) for f in s_keyframes]
durations = [
Duration(keyframes[d[0]], keyframes[d[1]], d[2], d[3]) for d in s_durations
]
return keyframes, durations
[docs]
class KeyframeWindow(Plugin):
def __init__(self):
self.reset()
@property
def scene(self):
return self.viewer.scene
@property
def keyframes(self):
return self.keyframe_editor.get_keyframes()
[docs]
def reset(self):
self.ui_window = None
self.popup = None
self.show_popup = False
self.keyframe_editor = None
self.key_frame_scenes = []
self.edited_duration = None
self._editor_file_name = None
self.current_frame = 0
self.total_frames = 32
[docs]
def close(self):
self.reset()
[docs]
def duration_name_change(self, text):
self.edited_duration.set_name(text.value)
[docs]
def duration_definition_change(self, text):
self.edited_duration.definition = text.value
[docs]
def open_in_editor(self, _):
if self._editor_file_name is None:
import tempfile
editor_file = tempfile.NamedTemporaryFile(
mode="w+",
prefix=self.edited_duration.name(),
suffix=".py",
delete=False,
)
editor_file.write(self.edited_duration.definition)
editor_file.close()
self._editor_file_name = editor_file.name
if platform.system() == "Linux":
os.system("gedit '{}' & disown".format(self._editor_file_name))
elif platform.system() == "Windows":
print(
"Please open and edit the file {} manually".format(
self._editor_file_name
)
)
[docs]
def notify_window_focus_change(self, focused):
if focused and self.edited_duration and self._editor_file_name:
with open(self._editor_file_name, "r") as f:
self.edited_duration.definition = f.read()
self.popup.get_children()[2].Value(self.edited_duration.definition)
[docs]
def build(self):
if self.scene is None:
self.ui_window = None
return
if self.ui_window:
return
self.popup = (
R.UIPopup()
.Label("Edit Duration")
.EscCallback(self.close_popup)
.append(
R.UISameLine().append(
R.UIButton().Label("Confirm").Callback(self.confirm_popup),
R.UIButton().Label("Cancel").Callback(self.cancel_popup),
R.UIButton().Label("Open in Editor").Callback(self.open_in_editor),
),
R.UIInputText().Label("Name").Callback(self.duration_name_change),
R.UIInputTextMultiline()
.Label("Definition")
.Callback(self.duration_definition_change),
)
)
self.keyframe_editor = (
R.UIKeyframeEditor(self.viewer.window.get_content_scale())
.BindCurrentFrame(self, "current_frame")
.BindTotalFrames(self, "total_frames")
.AddKeyframeCallback(self.add_keyframe)
.MoveKeyframeCallback(self.move_keyframe)
.AddDurationCallback(self.add_duration)
.DoubleClickKeyframeCallback(self.load_keyframe)
.DoubleClickDurationCallback(self.edit_duration)
.append(
R.UIButton().Label("Export").Callback(self.editor_export),
R.UIButton().Label("Import").Callback(self.editor_import),
)
)
self.ui_window = (
R.UIWindow().Label("Key Frame Editor").append(self.keyframe_editor)
)
[docs]
def add_keyframe(self, frame):
frame = Keyframe(SerializedScene(self.scene), frame)
self.keyframe_editor.add_keyframe(frame)
print(self.keyframes)
[docs]
def add_duration(self, frame0, frame1):
duration = Duration(frame0, frame1, "New Reward")
self.keyframe_editor.add_duration(duration)
[docs]
def move_keyframe(self, frame, time):
frame.set_frame(time)
[docs]
def load_keyframe(self, frame):
s_scene = frame.serialized_scene
s_scene.dump_state_into(self.scene)
[docs]
def edit_duration(self, duration):
self.edited_duration = duration
self.open_popup()
[docs]
def get_ui_windows(self):
self.build()
windows = []
if self.ui_window:
windows.append(self.ui_window)
if self.show_popup:
windows.append(self.popup)
return windows
[docs]
def get_editor_state(self):
keyframes = self.keyframe_editor.get_keyframes()
durations = self.keyframe_editor.get_durations()
return (
self.total_frames,
self.current_frame,
serialize_keyframes(keyframes, durations),
)
[docs]
def set_editor_state(self, state):
self.total_frames, self.current_frame, s = state
keyframes, durations = deserialize_keyframes(s)
self.keyframe_editor.set_state(keyframes, durations)
[docs]
def editor_export(self, _):
self._state = self.get_editor_state()
[docs]
def editor_import(self, _):
self.set_editor_state(self._state)