From b093e6b8afa78b7bfb8d4291b9a6a8f07a8c58f1 Mon Sep 17 00:00:00 2001 From: Manuel Kaufmann Date: Tue, 15 Aug 2023 22:59:32 +0200 Subject: [PATCH] - new function viewer.center_view_on_node to center on a node more easily in code. - added SpheresTrail and LinesTrail renderable --- aitviewer/headless.py | 8 ++++++-- aitviewer/renderables/lines.py | 23 +++++++++++++++++++++++ aitviewer/renderables/spheres.py | 16 ++++++++++++++++ aitviewer/scene/camera.py | 10 ++-------- aitviewer/viewer.py | 10 +++++++--- 5 files changed, 54 insertions(+), 13 deletions(-) diff --git a/aitviewer/headless.py b/aitviewer/headless.py index 2009411..8ae9e2b 100644 --- a/aitviewer/headless.py +++ b/aitviewer/headless.py @@ -51,8 +51,12 @@ def save_video(self, frame_dir=None, video_dir=None, output_fps=60, transparent= """ self._init_scene() self.export_video( - output_path=video_dir, frame_dir=frame_dir, animation=True, output_fps=output_fps, transparent=transparent, - **export_video_kwargs + output_path=video_dir, + frame_dir=frame_dir, + animation=True, + output_fps=output_fps, + transparent=transparent, + **export_video_kwargs, ) def save_frame(self, file_path, scale_factor: float = None): diff --git a/aitviewer/renderables/lines.py b/aitviewer/renderables/lines.py index 4e50648..21a4ad4 100644 --- a/aitviewer/renderables/lines.py +++ b/aitviewer/renderables/lines.py @@ -19,6 +19,7 @@ import trimesh from moderngl_window.opengl.vao import VAO +from aitviewer.renderables.spheres import SpheresTrail from aitviewer.scene.material import Material from aitviewer.scene.node import Node from aitviewer.shaders import ( @@ -588,3 +589,25 @@ def render(self, camera, **kwargs): def release(self): if self.is_renderable: self.vao.release() + + +class LinesTrail(Lines): + """A sequence of lines that leave a trail, i.e. the past lines keep being rendered.""" + + def __init__(self, lines, with_spheres=True, r_base=0.01, r_tip=None, color=(0.0, 0.0, 1.0, 1.0), **kwargs): + if "mode" not in kwargs: + kwargs["mode"] = "line_strip" + super().__init__(lines, r_base, r_tip, color, **kwargs) + self.n_frames = lines.shape[0] + + if with_spheres: + spheres_trail = SpheresTrail(lines, radius=r_base * 4.0, color=color) + self._add_node(spheres_trail, enabled=True, show_in_hierarchy=True) + + def on_frame_update(self): + self.n_lines = self.current_frame_id + super().on_frame_update() + + def make_renderable(self, ctx): + super().make_renderable(ctx) + self.on_frame_update() diff --git a/aitviewer/renderables/spheres.py b/aitviewer/renderables/spheres.py index b36d9ea..ceb44ac 100644 --- a/aitviewer/renderables/spheres.py +++ b/aitviewer/renderables/spheres.py @@ -278,3 +278,19 @@ def add_frames(self, positions): def remove_frames(self, frames): self.sphere_positions = np.delete(self.sphere_positions, frames, axis=0) self.redraw() + + +class SpheresTrail(Spheres): + """A sequence of spheres that leaves a trail, i.e. the past spheres keep being rendered.""" + + def __init__(self, positions, **kwargs): + super().__init__(positions, **kwargs) + self.n_frames = positions.shape[0] + + def on_frame_update(self): + self.n_spheres = self.current_frame_id + 1 + super().on_frame_update() + + def make_renderable(self, ctx): + super().make_renderable(ctx) + self.on_frame_update() diff --git a/aitviewer/scene/camera.py b/aitviewer/scene/camera.py index a4a784c..38a3cce 100644 --- a/aitviewer/scene/camera.py +++ b/aitviewer/scene/camera.py @@ -333,15 +333,10 @@ def show_path(self): all_oris[i] = self.rotation @ np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]]) path_spheres = RigidBodies(all_points, all_oris, radius=0.01, length=0.1, color=(0.92, 0.68, 0.2, 1.0)) + # Create lines only if there is more than one frame in the sequence. if self.n_frames > 1: - path_lines = Lines( - all_points, - color=(0, 0, 0, 1), - r_base=0.003, - mode="line_strip", - cast_shadow=False, - ) + path_lines = Lines(all_points, color=(0, 0, 0, 1), r_base=0.003, mode="line_strip", cast_shadow=False) else: path_lines = None @@ -351,7 +346,6 @@ def show_path(self): self.parent.add(path_spheres, show_in_hierarchy=False, enabled=self.enabled) if path_lines is not None: self.parent.add(path_lines, show_in_hierarchy=False, enabled=self.enabled) - self.path = (path_spheres, path_lines) self.current_frame_id = frame_id diff --git a/aitviewer/viewer.py b/aitviewer/viewer.py index 11197ad..817da72 100644 --- a/aitviewer/viewer.py +++ b/aitviewer/viewer.py @@ -1479,15 +1479,19 @@ def select_object(self, x: int, y: int): return False def center_view_on_selection(self): - if isinstance(self.scene.selected_object, Node): + self.center_view_on_node(self.scene.selected_object, with_animation=True) + + def center_view_on_node(self, node, with_animation=False): + if isinstance(node, Node): self.reset_camera() forward = self.scene.camera.forward - bounds = self.scene.selected_object.current_bounds + bounds = node.current_bounds diag = np.linalg.norm(bounds[:, 0] - bounds[:, 1]) dist = max(0.01, diag * 1.3) center = bounds.mean(-1) - self.scene.camera.move_with_animation(center - forward * dist, center) + anim_time = 0.25 if with_animation else 0.0 + self.scene.camera.move_with_animation(center - forward * dist, center, anim_time) def resize(self, width: int, height: int): self.window_size = (width, height)