Skip to content

Commit

Permalink
Added fast forward/Slow-Mo
Browse files Browse the repository at this point in the history 
Also added bugfix for number_of_steps
  • Loading branch information
@jpd-de
jpd-de authored Jun 10, 2024
1 parent 8925722 commit 21bacd8
Showing 1 changed file with 144 additions and 30 deletions.
174 changes: 144 additions & 30 deletions znvis/visualizer/visualizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,14 +95,19 @@ def __init__(
self.bounding_box = bounding_box() if bounding_box else None

if number_of_steps is None:
number_of_steps = len(particles[0].position)
len_list = []
for particle in particles:
if not particle.static:
len_list.append(len(particle.position))
number_of_steps = min(len_list)
self.number_of_steps = number_of_steps

self.output_folder = pathlib.Path(output_folder).resolve()
self.frame_folder = self.output_folder / "video_frames"
self.video_format = video_format
self.keep_frames = keep_frames
self.renderer = renderer
self.play_speed = 1

self.obj_folder = self.output_folder / "obj_files"

Expand Down Expand Up @@ -131,11 +136,15 @@ def _initialize_app(self):
self.vis.add_action("<<", self._update_particles_back)
self.vis.add_action("Play", self._continuous_trajectory)
self.vis.add_action(">>", self._update_particles)
self.vis.add_action(">>>", self._toggle_play_speed)
self.vis.add_action("Slow", self._toggle_slowmotion)
self.vis.add_action("Restart", self._restart_trajectory)
self.vis.add_action("Export Scene", self._export_scene)
self.vis.add_action("Screenshot", self._take_screenshot)
self.vis.add_action("Export Video", self._export_video)
self.vis.add_action("Export Mesh Trajectory", self._export_mesh_trajectory)


# Add the visualizer to the app.
self.app.add_window(self.vis)

Expand Down Expand Up @@ -272,18 +281,35 @@ def _take_screenshot(self, vis):
old_state = self.interrupt # get old state
self.interrupt = 0 # stop live feed if running.
mesh_dict = {}
mesh_center = []

if self.vector_field is not None:
for item in self.vector_field:
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

for item in self.particles:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
mesh_center.append(
item.mesh_list[self.counter]
.get_axis_aligned_bounding_box()
.get_center()
)
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

view_matrix = vis.scene.camera.get_view_matrix()

Expand Down Expand Up @@ -347,10 +373,12 @@ def _draw_particles(self, visualizer=None, initial: bool = False):
visualizer.add_geometry("Box", self.bounding_box)
else:
for i, item in enumerate(self.particles):
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)
if not item.static:
print(item.name)
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)


def _draw_vector_field(self, visualizer=None, initial: bool = False):
Expand All @@ -377,10 +405,11 @@ def _draw_vector_field(self, visualizer=None, initial: bool = False):
)
else:
for i, item in enumerate(self.vector_field):
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)
if not item.static:
visualizer.remove_geometry(item.name)
visualizer.add_geometry(
item.name, item.mesh_list[self.counter], item.mesh.o3d_material
)

def _continuous_trajectory(self, vis):
"""
Expand All @@ -396,6 +425,20 @@ def _continuous_trajectory(self, vis):
else:
threading.Thread(target=self._run_trajectory).start()

def _continuous_trajectory_backwards(self, vis):
"""
Button command for running the simulation in the visualizer backwards.
Parameters
----------
vis : visualizer
Object passed during the callback.
"""
if self.interrupt == 1:
self._pause_run(vis)
else:
threading.Thread(target=self._run_trajectory_backwards).start()

def _record_trajectory(self):
"""
Record the trajectory.
Expand All @@ -416,12 +459,34 @@ def save_callable():
"""
mesh_dict = {}

if self.vector_field is not None:
for item in self.vector_field:
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

for item in self.particles:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
if item.static:
mesh_dict[item.name] = {
"mesh": item.mesh_list[0],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}
else:
mesh_dict[item.name] = {
"mesh": item.mesh_list[self.counter],
"bsdf": item.mesh.material.mitsuba_bsdf,
"material": item.mesh.o3d_material,
}

view_matrix = self.vis.scene.camera.get_view_matrix()
self.renderer.render_mesh_objects(
Expand Down Expand Up @@ -473,10 +538,16 @@ def save_callable():
Function to be called on thread to save image.
"""
for i, item in enumerate(self.particles):
if i == 0:
mesh = item.mesh_list[self.counter]
if item.static:
if i == 0:
mesh = item.mesh_list[0]
else:
mesh += item.mesh_list[0]
else:
mesh += item.mesh_list[self.counter]
if i == 0:
mesh = item.mesh_list[self.counter]
else:
mesh += item.mesh_list[self.counter]

o3d.io.write_triangle_mesh(
(self.obj_folder / f"export_mesh_{self.counter}.ply").as_posix(), mesh
Expand Down Expand Up @@ -514,7 +585,7 @@ def _run_trajectory(self):
"""
self.interrupt = 1 # set global run state.
while self.counter < self.number_of_steps:
time.sleep(1 / self.frame_rate)
time.sleep(1 / (self.frame_rate * self.play_speed))
o3d.visualization.gui.Application.instance.post_to_main_thread(
self.vis, self._update_particles
)
Expand All @@ -524,6 +595,7 @@ def _run_trajectory(self):

self.interrupt = 0 # reset global state.


def _update_particles(self, visualizer=None, step: int = None):
"""
Update the positions of the particles.
Expand Down Expand Up @@ -570,12 +642,11 @@ def _update_particles_back(self, visualizer=None, step: int = None):
if visualizer is None:
visualizer = self.vis
if step is None:
if self.counter == self.number_of_steps - 1:
if self.counter == 0:
self.counter = self.number_of_steps - 1
else:
self.counter -= 1
step = self.counter

self._draw_particles(visualizer=visualizer) # draw the particles.

# draw the vector field if it exists.
Expand All @@ -584,6 +655,7 @@ def _update_particles_back(self, visualizer=None, step: int = None):

visualizer.post_redraw() # re-draw the window.


def _restart_trajectory(self, visualizer=None):
if visualizer is None:
visualizer = self.vis
Expand All @@ -597,7 +669,49 @@ def _restart_trajectory(self, visualizer=None):

visualizer.post_redraw() # re-draw the window.


def _toggle_play_speed(self, visualizer=None):
"""
Toggle the play speed from 1 to 2 to 4 to 8 and back to 1.
"""
if self.play_speed == 1:
self.play_speed = 2
elif self.play_speed == 2:
self.play_speed = 4
elif self.play_speed == 4:
self.play_speed = 8
else:
self.play_speed = 1

def _toggle_play_speed_back(self, visualizer=None):
"""
Toggle the play speed from 1 to 2 to 4 to 8 and back to 1.
"""

self._run_trajectory_backwards()
if self.play_speed == 1:
self.play_speed = 2
elif self.play_speed == 2:
self.play_speed = 4
elif self.play_speed == 4:
self.play_speed = 8
else:
self.play_speed = 1

def _toggle_slowmotion(self, visualizer=None):
"""
Toggle the play speed from 1 to 1/2 to 1/4 to 1/8 and back to 1
"""
if self.play_speed >= 1:
self.play_speed = 1/2
elif self.play_speed == 1/2:
self.play_speed = 1/4
elif self.play_speed == 1/4:
self.play_speed = 1/8
else:
self.play_speed = 1


def run_visualization(self):
"""
Run the visualization.
Expand Down

0 comments on commit 21bacd8

Please sign in to comment.