diff --git a/examples/biorbd/multi_models.py b/examples/biorbd/multi_models.py index 14d5d12..469a1cb 100644 --- a/examples/biorbd/multi_models.py +++ b/examples/biorbd/multi_models.py @@ -28,6 +28,7 @@ def main(): nb_seconds = 1 q0, t_span0 = building_some_q_and_t_span(nb_frames, nb_seconds) q1, t_span1 = building_some_q_and_t_span(20, 0.5) + time_offset = t_span0[-1] # loading biorbd model biorbd_model = BiorbdModel(biorbd_model_path) @@ -43,16 +44,17 @@ def main(): black_model = BiorbdModel(biorbd_model_path) black_model.options.mesh_color = (0, 0, 0) + black_model.options.show_gravity = True rerun_biorbd.add_animated_model(black_model, q0 + 0.2, phase=0, window="animation") rerun_biorbd.add_phase(t_span=t_span0, phase=0, window="split_animation") rerun_biorbd.add_animated_model(biorbd_model, q0 + 0.2, phase=0, window="split_animation") - rerun_biorbd.add_phase(t_span=t_span0[-1] + t_span1, phase=1) + rerun_biorbd.add_phase(t_span=time_offset + t_span1, phase=1) rerun_biorbd.add_animated_model(biorbd_model, q1, phase=1) - rerun_biorbd.add_phase(t_span=t_span0[-1] + t_span1, phase=1, window="split_animation") + rerun_biorbd.add_phase(t_span=time_offset + t_span1, phase=1, window="split_animation") rerun_biorbd.add_animated_model(biorbd_model, q1, phase=1, window="split_animation") markers = Markers(data=noisy_markers, channels=list(biorbd_model.marker_names)) diff --git a/pyorerun/abstract/abstract_class.py b/pyorerun/abstract/abstract_class.py index 6f8719d..c514a86 100644 --- a/pyorerun/abstract/abstract_class.py +++ b/pyorerun/abstract/abstract_class.py @@ -3,6 +3,16 @@ import numpy as np +class TimelessComponent(ABC): + @abstractmethod + def to_rerun(self): + pass + + @abstractmethod + def nb_components(self): + pass + + class Component(ABC): @abstractmethod def to_rerun(self, q: np.ndarray): diff --git a/pyorerun/biorbd_components/model_display_options.py b/pyorerun/biorbd_components/model_display_options.py index eed92e6..a5b39f9 100644 --- a/pyorerun/biorbd_components/model_display_options.py +++ b/pyorerun/biorbd_components/model_display_options.py @@ -19,3 +19,4 @@ class DisplayModelOptions: # NOTE : mesh_opacity doesnt exist in rerun yet # segment_frame_size: float = 0.1 not implemented yet mesh_color: tuple[int, int, int] = (255, 255, 255) + show_gravity: bool = False diff --git a/pyorerun/biorbd_components/model_interface.py b/pyorerun/biorbd_components/model_interface.py index 885de35..27effae 100644 --- a/pyorerun/biorbd_components/model_interface.py +++ b/pyorerun/biorbd_components/model_interface.py @@ -160,6 +160,10 @@ def q_ranges(self) -> tuple[tuple[float, float], ...]: q_ranges = [q_range for segment in self.model.segments() for q_range in segment.QRanges()] return tuple((q_range.min(), q_range.max()) for q_range in q_ranges) + @property + def gravity(self) -> np.ndarray: + return self.model.getGravity().to_array() + class BiorbdModel(BiorbdModelNoMesh): """ diff --git a/pyorerun/multi_phase_rerun.py b/pyorerun/multi_phase_rerun.py index 195be6c..09a3d8e 100644 --- a/pyorerun/multi_phase_rerun.py +++ b/pyorerun/multi_phase_rerun.py @@ -1,5 +1,6 @@ import numpy as np import rerun as rr # NOTE: `rerun`, not `rerun-sdk`! +import rerun.blueprint as rrb from pyomeca import Markers as PyoMarkers from .biorbd_components.model_interface import BiorbdModel @@ -12,7 +13,7 @@ class MultiPhaseRerun: """ def __init__(self) -> None: - self.rerun_biorbd_phases: list[dict] = [] + self.rerun_biorbd_phases: list[dict[str, PhaseRerun], ...] = [] def add_phase(self, t_span: np.ndarray, phase: int = 0, window: str = "animation") -> None: @@ -48,6 +49,12 @@ def all_windows(self) -> list[str]: def rerun(self, server_name: str = "multi_phase_animation") -> None: rr.init(server_name, spawn=True) for i, phase in enumerate(self.rerun_biorbd_phases): - for j, rr_phase in enumerate(phase.values()): + for j, (window, rr_phase) in enumerate(phase.items()): + + rrb.Spatial3DView( + origin="/", + contents=f"{window}/**", + ) + more_phases_after_this_one = i < self.nb_phase - 1 rr_phase.rerun(init=False, clear_last_node=more_phases_after_this_one) diff --git a/pyorerun/phase_rerun.py b/pyorerun/phase_rerun.py index e0fc4b6..ef30879 100644 --- a/pyorerun/phase_rerun.py +++ b/pyorerun/phase_rerun.py @@ -5,6 +5,8 @@ from .abstract.q import QProperties from .biorbd_components.model_interface import BiorbdModel from .biorbd_phase import BiorbdRerunPhase +from .timeless.gravity import Gravity +from .timeless_components import TimelessRerunPhase from .xp_components.markers import MarkersXp from .xp_components.timeseries_q import TimeSeriesQ from .xp_phase import XpRerunPhase @@ -13,6 +15,22 @@ class PhaseRerun: """ A class to animate a biorbd model in rerun. + + Attributes + ---------- + phase : int + The phase number of the animation. + name : str + The name of the animation. + t_span : np.ndarray + The time span of the animation, such as the time instant of each frame. + biorbd_models : BiorbdRerunPhase + The biorbd models to animate. + xp_data : XpRerunPhase + The experimental data to display. + timeless_components : list + The components to display at the begin of the phase but stay until the end of this phase. + This not a true timeless in the sens of rerun, as if a new phase is created, the timeless components will be cleared. """ def __init__(self, t_span: np.ndarray, phase: int = 0, window: str = None): @@ -34,6 +52,7 @@ def __init__(self, t_span: np.ndarray, phase: int = 0, window: str = None): self.biorbd_models = BiorbdRerunPhase(name=self.name, phase=phase) self.xp_data = XpRerunPhase(name=self.name, phase=phase) + self.timeless_components = TimelessRerunPhase(name=self.name, phase=phase) def add_animated_model( self, biomod: BiorbdModel, q: np.ndarray, tracked_markers: PyoMarkers = None, display_q: bool = False @@ -73,6 +92,10 @@ def add_animated_model( ranges=biomod.q_ranges, dof_names=biomod.dof_names, ) + if biomod.options.show_gravity: + self.timeless_components.add_component( + Gravity(name=f"{self.name}/{self.biorbd_models.nb_models}_{biomod.name}", vector=biomod.gravity) + ) def __add_tracked_markers(self, biomod: BiorbdModel, tracked_markers: PyoMarkers) -> None: """Add the tracked markers to the phase.""" @@ -138,11 +161,21 @@ def rerun(self, name: str = "animation_phase", init: bool = True, clear_last_nod if init: rr.init(f"{name}_{self.phase}", spawn=True) - for frame, t in enumerate(self.t_span): + frame = 0 + rr.set_time_seconds("stable_time", self.t_span[frame]) + self.timeless_components.to_rerun() + self.biorbd_models.to_rerun(frame) + self.xp_data.to_rerun(frame) + + for frame, t in enumerate(self.t_span[1:]): rr.set_time_seconds("stable_time", t) - self.biorbd_models.to_rerun(frame) - self.xp_data.to_rerun(frame) + self.biorbd_models.to_rerun(frame + 1) + self.xp_data.to_rerun(frame + 1) if clear_last_node: - for component in [*self.biorbd_models.component_names, *self.xp_data.component_names]: + for component in [ + *self.biorbd_models.component_names, + *self.xp_data.component_names, + *self.timeless_components.component_names, + ]: rr.log(component, rr.Clear(recursive=False)) diff --git a/pyorerun/timeless/gravity.py b/pyorerun/timeless/gravity.py new file mode 100644 index 0000000..6282c3c --- /dev/null +++ b/pyorerun/timeless/gravity.py @@ -0,0 +1,21 @@ +import numpy as np +import rerun as rr + +from ..abstract.abstract_class import TimelessComponent + + +class Gravity(TimelessComponent): + + def __init__(self, name, vector: np.ndarray): + self.name = name + "/gravity" + self.vector = vector / 20 + + @property + def nb_components(self): + return 1 + + def to_rerun(self) -> None: + rr.log( + self.name, + rr.Arrows3D(origins=np.zeros(3), vectors=self.vector, colors=np.array([255, 255, 255])), + ) diff --git a/pyorerun/timeless_components.py b/pyorerun/timeless_components.py new file mode 100644 index 0000000..d680532 --- /dev/null +++ b/pyorerun/timeless_components.py @@ -0,0 +1,19 @@ +from typing import Any + + +class TimelessRerunPhase: + def __init__(self, name, phase: int): + self.name = name + self.phase = phase + self.timeless_components = [] + + def add_component(self, component: Any): + self.timeless_components.append(component) + + def to_rerun(self): + for data in self.timeless_components: + data.to_rerun() + + @property + def component_names(self) -> list[str]: + return [data.name for data in self.timeless_components]