[RTR] Muscles changing color with activation

environment.yml

@@ -7,9 +7,11 @@ dependencies:
 - ezc3d
 - rerun-sdk=0.21.0
 - numpy
+- matplotlib
 - biorbd>=1.10.5
 - trimesh
 - tk
 - imageio
 - imageio-ffmpeg
 - opensim
+

examples/biorbd/gait_reconstruction.py

@@ -1,8 +1,7 @@
 import pickle
+import numpy as np
 
-from pyorerun.pyomarkers import Pyomarkers as PyoMarkers
-
-from pyorerun import BiorbdModel, PhaseRerun
+from pyorerun import BiorbdModel, PhaseRerun, PyoMuscles, PyoMarkers
 
 
 def main():
@@ -34,6 +33,19 @@ def main():
     # Add experimental markers
     pyomarkers = PyoMarkers(data=markers, channels=list(model.marker_names), show_labels=False)
 
+    # Add experimental emg
+    nb_muscles = model.nb_muscles
+    nb_frames = q.shape[1]
+    fake_emg = np.ones((nb_muscles, nb_frames))  # Fake EMG data for demonstration
+    for i_muscle in range(nb_muscles):
+        fake_emg[i_muscle, :] = np.linspace(0.01, 1, nb_frames)
+    pyoemg = PyoMuscles(
+        data=fake_emg,
+        muscle_names=list(model.muscle_names),
+        mvc=np.ones((nb_muscles,)),  # Fake MVC values
+        colormap="viridis",
+    )
+
     # Add force plates to the animation
     viz.add_force_plate(num=0, corners=force_plate_1_corners)
     viz.add_force_plate(num=1, corners=force_plate_2_corners)
@@ -52,7 +64,7 @@ def main():
     viz.add_animated_model(
         model, q
     )  # This line is just to test the model without markers (but is not necessary for the example to work)
-    viz.add_animated_model(model, q, tracked_markers=pyomarkers)
+    viz.add_animated_model(model, q, tracked_markers=pyomarkers, muscle_activations_intensity=pyoemg)
 
     # Play
     viz.rerun("Experimental data with kinematics reconstruction")

examples/biorbd/marker_tracking.py

@@ -1,7 +1,6 @@
 import numpy as np
 
-from pyorerun import BiorbdModel, PhaseRerun
-from pyorerun.pyomarkers import Pyomarkers
+from pyorerun import BiorbdModel, PhaseRerun, PyoMarkers
 
 
 def main():
@@ -26,7 +25,7 @@ def main():
 
     # running the animation
     rerun_biorbd = PhaseRerun(t_span)
-    markers = Pyomarkers(data=noisy_markers, channels=list(biorbd_model.marker_names))
+    markers = PyoMarkers(data=noisy_markers, channels=list(biorbd_model.marker_names))
     rerun_biorbd.add_animated_model(biorbd_model, q, tracked_markers=markers)
 
     # rerun_biorbd.add_xp_markers(

examples/biorbd/multi_models.py

@@ -6,9 +6,8 @@
 
 import numpy as np
 from numpy import random
-from pyorerun.pyomarkers import Pyomarkers as Markers
 
-from pyorerun import BiorbdModel, MultiPhaseRerun
+from pyorerun import BiorbdModel, MultiPhaseRerun, PyoMarkers
 
 
 def building_some_q_and_t_span(nb_frames: int, nb_seconds: int) -> tuple[np.ndarray, np.ndarray]:
@@ -59,7 +58,7 @@ def main():
     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))
+    markers = PyoMarkers(data=noisy_markers, channels=list(biorbd_model.marker_names))
     rerun_biorbd.add_xp_markers(
         name="noisy_markers",
         markers=markers,

examples/biorbd/no_mesh.py

@@ -1,7 +1,6 @@
 import numpy as np
 
-from pyorerun import BiorbdModelNoMesh, PhaseRerun
-from pyorerun.pyomarkers import Pyomarkers as PyoMarkers
+from pyorerun import BiorbdModelNoMesh, PhaseRerun, PyoMarkers
 
 
 def main():

pyorerun/__init__.py

@@ -15,7 +15,8 @@
 
 # from .biorbd_phase import rr_biorbd as rrbiorbd, RerunBiorbdPhase
 from .phase_rerun import PhaseRerun
-from .pyomarkers import Pyomarkers
+from .pyomarkers import PyoMarkers
+from .pyoemg import PyoMuscles
 from .rrbiomod import rr_biorbd as animate
 from .rrc3d import rrc3d as c3d
 from .xp_components.timeseries_q import OsimTimeSeries

pyorerun/abstract/linestrip.py

@@ -91,7 +91,7 @@ def radius_to_rerun(self) -> np.ndarray:
         else:
             return self.radius
 
-    def color_to_rerun(self) -> np.ndarray:
+    def color_to_rerun(self, nb_frames: int) -> np.ndarray:
         """
         Returns a numpy array with the color of each line.
 
@@ -100,10 +100,14 @@ def color_to_rerun(self) -> np.ndarray:
         np.ndarray
             A numpy array with the color of each line.
         """
-        if self.color.ndim == 1:
-            return np.tile(self.color, (self.nb_strips, 1))
+        nb_strips = len(self.strip_names)
+        if len(self.color.shape) == 3:
+            colors = [self.color[s, f, :] for f in range(nb_frames) for s in range(nb_strips)]
+        elif len(self.color.shape) == 2:
+            colors = [self.color[s, :, :] for _ in range(nb_frames) for s in range(nb_strips)]
         else:
-            return self.color
+            colors = [self.color for _ in range(nb_frames * nb_strips)]
+        return colors
 
     def show_labels_to_rerun(self) -> list[bool]:
         """

pyorerun/live_animation.py

@@ -7,6 +7,7 @@
 from .model_interfaces.biorbd_model_interface import BiorbdModel
 from .model_interfaces.osim_model_interface import OsimModel
 from .model_components.model_updapter import ModelUpdater
+from .model_components.model_display_options import DisplayModelOptions
 
 
 class LiveModelAnimation:
@@ -53,6 +54,7 @@ def __init__(self, model_path: str, with_q_charts: bool = False):
         self.dof_slider_values = []
         self.update_functions = []
         self.with_q_charts = with_q_charts
+        self.options = DisplayModelOptions()
 
     def update_viewer(self, event, dof_index: int):
         the_dof_idx, the_value = self.fetch_and_update_slider_value(event, dof_index)

pyorerun/model_components/ligaments.py

@@ -28,7 +28,7 @@ def to_component(self, q: np.ndarray) -> rr.LineStrips3D:
         return rr.LineStrips3D(
             strips=self.update_callable(q),
             radii=self.properties.radius_to_rerun(),
-            colors=self.properties.color_to_rerun(),
+            colors=self.properties.color_to_rerun(1),
             labels=self.properties.strip_names,
             show_labels=self.properties.show_labels_to_rerun(),
         )
@@ -62,7 +62,7 @@ def to_chunk(self, q: np.ndarray) -> dict[str, list]:
         #     for s in range(self.nb_strips)
         # }
         # lets try a more advanced approach
-        colors = [self.properties.color for _ in range(nb_frames * self.nb_strips)]
+        colors = self.properties.color_to_rerun(nb_frames)
         radii = [self.properties.radius for _ in range(nb_frames * self.nb_strips)]
         labels = [self.properties.strip_names[s] for _ in range(nb_frames) for s in range(self.nb_strips)]
         partition = [self.nb_strips for _ in range(nb_frames)]
@@ -125,7 +125,7 @@ def to_rerun(self, q: np.ndarray = None, markers: np.ndarray = None) -> None:
             rr.LineStrips3D(
                 strips=self.line_strips(q, markers),
                 radii=self.properties.radius_to_rerun(),
-                colors=self.properties.color_to_rerun(),
+                colors=self.properties.color_to_rerun(1),
                 # labels=self.properties.strip_names,
             ),
         )
@@ -180,7 +180,7 @@ def __init__(self, name, properties: LineStripProperties, strips: np.ndarray, tr
         self.rerun_mesh = rr.LineStrips3D(
             strips=strips,
             radii=self.properties.radius_to_rerun(),
-            colors=self.properties.color_to_rerun(),
+            colors=self.properties.color_to_rerun(1),
         )
 
     def initialize(self):

pyorerun/model_components/model_updapter.py

@@ -16,7 +16,12 @@
 
 
 class ModelUpdater(Components):
-    def __init__(self, name, model: BiorbdModelNoMesh | BiorbdModel | OsimModelNoMesh | OsimModel | AbstractModel):
+    def __init__(
+        self,
+        name: str,
+        model: BiorbdModelNoMesh | BiorbdModel | OsimModelNoMesh | OsimModel | AbstractModel,
+        muscle_colors: np.ndarray = None,
+    ):
         self.name = name
         self.model = model
         self.markers = self.create_markers_updater()
@@ -25,7 +30,7 @@ def __init__(self, name, model: BiorbdModelNoMesh | BiorbdModel | OsimModelNoMes
         self.rigid_contacts = self.create_rigid_contacts_updater()
         self.ligaments = self.create_ligaments_updater()
         self.segments = self.create_segments_updater()
-        self.muscles = self.create_muscles_updater()
+        self.muscles = self.create_muscles_updater(muscle_colors)
 
     @classmethod
     def from_file(cls, model_path: str, options: DisplayModelOptions = None):
@@ -186,14 +191,15 @@ def create_segments_updater(self):
             segments.append(SegmentUpdater(name=segment_name, transform_callable=transform_callable, meshes=meshes))
         return segments
 
-    def create_muscles_updater(self):
+    def create_muscles_updater(self, muscle_colors: np.ndarray = None):
         if self.model.nb_muscles == 0:
             return EmptyUpdater(self.name + "/muscles")
+        colors = muscle_colors if muscle_colors is not None else np.array(self.model.options.muscles_color)
         return MusclesUpdater(
             self.name,
             properties=LineStripProperties(
                 strip_names=self.model.muscle_names,
-                color=np.array(self.model.options.muscles_color),
+                color=colors,
                 radius=self.model.options.muscles_radius,
                 show_labels=self.model.options.show_muscle_labels,
             ),

pyorerun/model_interfaces/biorbd_model_interface.py

@@ -107,8 +107,12 @@ def segment_homogeneous_matrices_in_global(self, q: np.ndarray, segment_index: i
         Returns a biorbd object containing the roto-translation matrix of the segment in the global reference frame.
         This is useful if you want to interact with biorbd directly later on.
         """
-        rt_matrix = self.model.globalJCS(GeneralizedCoordinates(q), segment_index)
-        return rt_matrix.to_array()
+        if np.sum(np.isnan(q)) != 0:
+            # If q contains NaN, return an identity matrix as biorbd will throw an error otherwise
+            rt_matrix = np.identity(4)
+        else:
+            rt_matrix = self.model.globalJCS(GeneralizedCoordinates(q), segment_index).to_array()
+        return rt_matrix
 
     def markers(self, q: np.ndarray) -> np.ndarray:
         """
@@ -285,9 +289,18 @@ def mesh_homogenous_matrices_in_global(self, q: np.ndarray, segment_index: int,
         """
         Returns a list of homogeneous matrices of the mesh in the global reference frame
         """
-        mesh_rt = (
-            super(BiorbdModel, self).segments[segment_index].segment.characteristics().mesh().getRotation().to_array()
-        )
-        # mesh_rt = self.segments[segment_index].segment.characteristics().mesh().getRotation().to_array()
-        segment_rt = self.segment_homogeneous_matrices_in_global(q, segment_index=segment_index)
-        return segment_rt @ mesh_rt
+        if np.sum(np.isnan(q)) != 0:
+            # If q contains NaN, return an identity matrix as biorbd will throw an error otherwise
+            return np.identity(4)
+        else:
+            mesh_rt = (
+                super(BiorbdModel, self)
+                .segments[segment_index]
+                .segment.characteristics()
+                .mesh()
+                .getRotation()
+                .to_array()
+            )
+            # mesh_rt = self.segments[segment_index].segment.characteristics().mesh().getRotation().to_array()
+            segment_rt = self.segment_homogeneous_matrices_in_global(q, segment_index=segment_index)
+            return segment_rt @ mesh_rt

pyorerun/model_phase.py

@@ -16,6 +16,7 @@ def __init__(self, name, phase: int = 0):
         self.models = []
         self.rerun_models = []
         self.q = []
+        self.muscle_colors = []
         self.tracked_markers = []
         self.rerun_links = []
 
@@ -27,18 +28,22 @@ def _rerun_links_without_none(self):
     def _model_links_index_without_none(self):
         return [i for i, rr_link in enumerate(self.rerun_links) if rr_link is not None]
 
-    def add_animated_model(self, model: AbstractModel, q: np.ndarray, tracked_markers: np.ndarray = None):
+    def add_animated_model(
+        self, model: AbstractModel, q: np.ndarray, tracked_markers: np.ndarray = None, muscle_colors: np.ndarray = None
+    ):
         self.models.append(model)
-        self.rerun_models.append(ModelUpdater(name=f"{self.name}/{self.nb_models}_{model.name}", model=model))
+        self.rerun_models.append(
+            ModelUpdater(name=f"{self.name}/{self.nb_models}_{model.name}", model=model, muscle_colors=muscle_colors)
+        )
         self.q.append(q)
 
         self.tracked_markers.append(tracked_markers if tracked_markers is not None else None)
-        updater = (
+        marker_updater = (
             ModelMarkerLinksUpdater(name=f"{self.name}/{self.nb_models}_{model.name}", model=model)
             if tracked_markers is not None
             else None
         )
-        self.rerun_links.append(updater)
+        self.rerun_links.append(marker_updater)
 
     def to_rerun(self, frame: int):
         self.to_rerun_models(frame)

pyorerun/multi_phase_rerun.py

@@ -1,7 +1,7 @@
 import numpy as np
 import rerun as rr  # NOTE: `rerun`, not `rerun-sdk`!
 import rerun.blueprint as rrb
-from .pyomarkers import Pyomarkers as PyoMarkers
+from .pyomarkers import PyoMarkers
 
 from .model_interfaces import AbstractModel
 from .phase_rerun import PhaseRerun

pyorerun/phase_rerun.py

@@ -1,6 +1,7 @@
 import numpy as np
 import rerun as rr
-from .pyomarkers import Pyomarkers as PyoMarkers
+from .pyomarkers import PyoMarkers
+from .pyoemg import PyoMuscles
 
 from .abstract.q import QProperties
 from .model_interfaces import AbstractModel
@@ -58,7 +59,8 @@ def add_animated_model(
         self,
         model: AbstractModel,
         q: np.ndarray,
-        tracked_markers: PyoMarkers = None,
+        tracked_markers: PyoMarkers | np.ndarray = None,
+        muscle_activations_intensity: PyoMuscles | np.ndarray = None,
         display_q: bool = False,
     ) -> None:
         """
@@ -72,6 +74,8 @@ def add_animated_model(
             The generalized coordinates of the model.
         tracked_markers: PyoMarkers
             The markers to display, and sets a link between the model markers and the tracked markers.
+        muscle_activations_intensity: PyoMuscles
+            The muscle activation level to display.
         display_q: bool
             Whether to display the generalized coordinates q in charts.
         """
@@ -83,14 +87,22 @@ def add_animated_model(
                 f"Current shapes are q: {q.shape[1]} and tspan: {self.t_span.shape}."
             )
 
-        if tracked_markers is None:
-            self.models.add_animated_model(model, q)
-        else:
-            if isinstance(tracked_markers, np.ndarray):
-                tracked_markers = PyoMarkers(tracked_markers, channels=model.marker_names)
-            self.models.add_animated_model(model, q, tracked_markers.to_numpy()[:3, :, :])
+        if isinstance(tracked_markers, np.ndarray):
+            tracked_markers = PyoMarkers(tracked_markers, channels=model.marker_names)
+
+        muscle_colors = None
+        if isinstance(muscle_activations_intensity, np.ndarray):
+            muscle_colors = PyoMuscles(muscle_activations_intensity, muscle_names=model.muscle_names).to_colors()
+        elif isinstance(muscle_activations_intensity, PyoMuscles):
+            muscle_colors = muscle_activations_intensity.to_colors()
+
+        if tracked_markers is not None:
             self.__add_tracked_markers(model, tracked_markers)
 
+        if tracked_markers is not None:
+            tracked_markers = tracked_markers.to_numpy()[:3, :, :]
+        self.models.add_animated_model(model, q, tracked_markers, muscle_colors)
+
         if display_q:
             self.add_q(
                 f"{model.name}_q",