Merge pull request #90 from pariterre/master

Added the distance between markers and experimental markers

Author: pariterre

bioviz/__init__.py

@@ -146,7 +146,6 @@ def __init__(
             soft_contacts_size=soft_contacts_size,
             soft_contacts_color=soft_contacts_color,
         )
-        self.vtk_model_markers: VtkModel | None = None
         self.is_executing = False
         self.animation_warning_already_shown = False
 
@@ -161,6 +160,7 @@ def __init__(
         self.show_experimental_markers = False
         self.experimental_markers = None
         self.experimental_markers_color = experimental_markers_color
+        self.virtual_to_experimental_markers_indices = None
         self.show_experimental_forces = False
         self.experimental_forces = None
         self.segment_forces = []
@@ -348,6 +348,10 @@ def __init__(
 
         # Update everything at the position Q=0
         self.set_q(self.Q)
+        if self.show_floor:
+            self.__set_floor()
+        if self.show_gravity_vector:
+            self.__set_gravity_vector()
 
     def reset_q(self):
         self.Q = np.zeros(self.Q.shape)
@@ -386,8 +390,6 @@ def set_q(self, Q, refresh_window=True):
         self.__set_rt_from_q()
         self.__set_meshes_from_q()
         self.__set_global_center_of_mass_from_q()
-        self.__set_gravity_vector()
-        self.__set_floor()
         self.__set_segments_center_of_mass_from_q()
         self.__set_markers_from_q()
         self.__set_contacts_from_q()
@@ -1032,12 +1034,32 @@ def __load_experimental_data_from_button(self):
     def load_c3d(self, data, auto_start=True, ignore_animation_warning=True):
         self.load_experimental_markers(data, auto_start, ignore_animation_warning)
 
-    def load_experimental_markers(self, data, auto_start=True, ignore_animation_warning=True):
+    def load_experimental_markers(
+        self,
+        data,
+        auto_start=True,
+        ignore_animation_warning=True,
+        experimental_markers_mapping_to_virtual: list[int, ...] = None,
+    ):
         if isinstance(data, str):
             self.experimental_markers = Markers.from_c3d(data)
             if self.experimental_markers.units == "mm":
                 self.experimental_markers = self.experimental_markers * 0.001
 
+            # Try to find a correspondence between the loaded experimental markers and the model
+            self.virtual_to_experimental_markers_indices = experimental_markers_mapping_to_virtual
+            if self.virtual_to_experimental_markers_indices is None:
+                try:
+                    virtual_marker_names = [n.to_string() for n in self.model.markerNames()]
+                    exp_marker_names = list(self.experimental_markers.channel.data)
+                    self.virtual_to_experimental_markers_indices = [
+                        virtual_marker_names.index(name) if name in virtual_marker_names else None
+                        for name in exp_marker_names
+                    ]
+                except ValueError:
+                    # Did not find direct correspondence
+                    pass
+
             self.c3d_file_name = data
             self.radio_c3d_editor_model.setEnabled(True)
 
@@ -1050,11 +1072,6 @@ def load_experimental_markers(self, data, auto_start=True, ignore_animation_warn
                 f"Allowed type are numpy array (3xNxT), data array (3xNxT) or .c3d file (str)."
             )
 
-        if not self.vtk_model_markers:
-            self.vtk_model_markers = VtkModel(
-                self.vtk_window, markers_color=self.experimental_markers_color, markers_size=self.vtk_markers_size
-            )
-
         self.__set_movement_slider()
         self.show_experimental_markers = True
 
@@ -1110,7 +1127,11 @@ def __set_experimental_markers_from_frame(self):
 
         t_slider = self.movement_slider[0].value() - 1
         t = t_slider if t_slider < self.experimental_markers.shape[2] else self.experimental_markers.shape[2] - 1
-        self.vtk_model_markers.update_markers(self.experimental_markers[:, :, t : t + 1].isel(time=[0]))
+        self.vtk_model.update_experimental_markers(
+            self.experimental_markers[:, :, t : t + 1].isel(time=[0]),
+            with_link=True,
+            virtual_to_experimental_markers_indices=self.virtual_to_experimental_markers_indices,
+        )
 
     def __set_experimental_forces_from_frame(self):
         if not self.show_experimental_forces:

bioviz/_version.py

@@ -1,6 +1,6 @@
 from packaging.version import parse as parse_version
 
-__version__ = "2.3.0"
+__version__ = "2.3.2"
 
 
 def check_version(tool_to_compare, min_version, max_version):

bioviz/biorbd_vtk.py

@@ -1,6 +1,7 @@
 """
 Visualization toolkit in pyomeca
 """
+from dataclasses import dataclass
 import os
 import time
 import sys
@@ -198,13 +199,25 @@ def record(self, finish=False, button_to_block=(), file_name=None):
                 self.setMaximumSize(self.maximum_size)
 
 
+@dataclass
+class _MarkerInternal:
+    data: Markers
+    size: float
+    color: tuple[float, float, float]
+    opacity: float
+    actors: list[QVTKRenderWindowInteractor, ...]
+
+
 class VtkModel(QtWidgets.QWidget):
     def __init__(
         self,
         parent,
         markers_size=0.010,
         markers_color=(1, 1, 1),
         markers_opacity=1.0,
+        experimental_markers_size=0.010,
+        experimental_markers_color=(1, 1, 1),
+        experimental_markers_opacity=1.0,
         contacts_color=(0, 1, 0),
         contacts_size=0.01,
         contacts_opacity=1.0,
@@ -258,11 +271,19 @@ def __init__(
         self.setAutoFillBackground(True)
         self.setPalette(palette)
 
-        self.markers = Markers()
-        self.markers_size = markers_size
-        self.markers_color = markers_color
-        self.markers_opacity = markers_opacity
-        self.markers_actors = list()
+        self.markers = {
+            "model": _MarkerInternal(
+                data=Markers(), color=markers_color, size=markers_size, opacity=markers_opacity, actors=list()
+            ),
+            "experimental": _MarkerInternal(
+                data=Markers(),
+                color=experimental_markers_color,
+                size=experimental_markers_size,
+                opacity=experimental_markers_opacity,
+                actors=list(),
+            ),
+        }
+        self.markers_link_actors: list[QVTKRenderWindowInteractor, ...] = list()
 
         self.contacts = Markers()
         self.contacts_size = contacts_size
@@ -343,21 +364,53 @@ def set_markers_color(self, markers_color):
         markers_color : tuple(int)
             Color the markers should be drawn (1 is max brightness)
         """
-        self.markers_color = markers_color
-        self.update_markers(self.markers)
+        self._set_markers_color(markers_color, key="model")
+
+    def set_experimental_markers_color(self, markers_color):
+        """
+        Dynamically change the color of the markers
+        Parameters
+        ----------
+        markers_color : tuple(int)
+            Color the markers should be drawn (1 is max brightness)
+        """
+        self._set_markers_color(markers_color, key="experimental")
+
+    def _set_markers_color(self, markers_color, key):
+        """
+        Dynamically change the color of the markers
+        Parameters
+        ----------
+        markers_color : tuple(int)
+            Color the markers should be drawn (1 is max brightness)
+        """
+        self.markers_color[key] = markers_color
+        self._update_markers(self.markers, key)
 
     def set_markers_size(self, markers_size):
+        self._set_markers_size(markers_size, "model")
+
+    def set_experimental_markers_size(self, markers_size):
+        self._set_markers_size(markers_size, "experimental")
+
+    def _set_markers_size(self, markers_size, key):
         """
         Dynamically change the size of the markers
         Parameters
         ----------
         markers_size : float
             Size the markers should be drawn
         """
-        self.markers_size = markers_size
-        self.update_markers(self.markers)
+        self.markers_size[key] = markers_size
+        self._update_markers(self.markers, key)
 
     def set_markers_opacity(self, markers_opacity):
+        self._set_markers_opacity(markers_opacity, "model")
+
+    def set_experimental_markers_opacity(self, markers_opacity):
+        self._set_markers_opacity(markers_opacity, "experimental")
+
+    def _set_markers_opacity(self, markers_opacity, key):
         """
         Dynamically change the opacity of the markers
         Parameters
@@ -368,10 +421,10 @@ def set_markers_opacity(self, markers_opacity):
         -------
 
         """
-        self.markers_opacity = markers_opacity
-        self.update_markers(self.markers)
+        self.markers_opacity[key] = markers_opacity
+        self._update_markers(self.markers, key)
 
-    def new_marker_set(self, markers):
+    def _new_marker_set(self, markers, key):
         """
         Define a new marker set. This function must be called each time the number of markers change
         Parameters
@@ -382,28 +435,52 @@ def new_marker_set(self, markers):
         """
         if len(markers.shape) > 2 and markers.shape[2] > 1:
             raise IndexError("Markers should be from one frame only")
-        self.markers = markers
+        self.markers[key].data = markers
 
         # Remove previous actors from the scene
-        for actor in self.markers_actors:
+        for actor in self.markers[key].actors:
             self.parent_window.ren.RemoveActor(actor)
-        self.markers_actors = list()
+        self.markers[key].actors = list()
 
         # Create the geometry of a point (the coordinate) points = vtk.vtkPoints()
         for i in range(markers.channel.size):
             # Create a mapper
             mapper = vtkPolyDataMapper()
 
             # Create an actor
-            self.markers_actors.append(vtkActor())
-            self.markers_actors[i].SetMapper(mapper)
+            self.markers[key].actors.append(vtkActor())
+            self.markers[key].actors[i].SetMapper(mapper)
 
-            self.parent_window.ren.AddActor(self.markers_actors[i])
+            self.parent_window.ren.AddActor(self.markers[key].actors[i])
 
         # Update marker position
-        self.update_markers(self.markers)
+        self._update_markers(self.markers[key].data, key)
 
     def update_markers(self, markers):
+        self._update_markers(markers, "model")
+
+    def update_experimental_markers(
+        self, markers, with_link: bool = True, virtual_to_experimental_markers_indices: list[int, ...] = None
+    ):
+        """
+        Update position of the experimental markers on the screen (but do not repaint)
+        Parameters
+        ----------
+        markers : Markers3d
+            One frame of markers
+        with_link : bool
+            If links to the virtual markers should be added (implies that coll to update_markers was done prior
+        virtual_to_experimental_markers_indices : list[int, ...]
+            A list that links the virtual markers indices to the experimental marker indices. If left None,
+            it is assumed to be a one to one (0=>0, 1=>1, ..., n_virtual_markers=>n_experimental_markers)
+        """
+        self._update_markers(markers, "experimental")
+        if with_link:
+            if virtual_to_experimental_markers_indices is None:
+                virtual_to_experimental_markers_indices = tuple(range(self.markers["model"].data.channel.size))
+            self._update_experimental_marker_link(virtual_to_experimental_markers_indices)
+
+    def _update_markers(self, markers, key):
         """
         Update position of the markers on the screen (but do not repaint)
         Parameters
@@ -415,22 +492,90 @@ def update_markers(self, markers):
 
         if len(markers.shape) > 2 and markers.shape[2] > 1:
             raise IndexError("Markers should be from one frame only")
-        if markers.channel.size != self.markers.channel.size:
-            self.new_marker_set(markers)
+        if markers.channel.size != self.markers[key].data.channel.size:
+            self._new_marker_set(markers, key)
             return  # Prevent calling update_markers recursively
-        self.markers = markers
+        self.markers[key].data = markers
         markers = np.array(markers)
 
-        for i, actor in enumerate(self.markers_actors):
+        for i, actor in enumerate(self.markers[key].actors):
             # mapper = actors.GetNextActor().GetMapper()
             mapper = actor.GetMapper()
-            self.markers_actors[i].GetProperty().SetColor(self.markers_color)
-            self.markers_actors[i].GetProperty().SetOpacity(self.markers_opacity)
+            self.markers[key].actors[i].GetProperty().SetColor(self.markers[key].color)
+            self.markers[key].actors[i].GetProperty().SetOpacity(self.markers[key].opacity)
             source = vtkSphereSource()
             source.SetCenter(markers[0:3, i])
-            source.SetRadius(self.markers_size)
+            source.SetRadius(self.markers[key].size)
             mapper.SetInputConnection(source.GetOutputPort())
 
+    def _new_experimental_marker_link(self, virtual_to_experimental_markers_indices):
+        """
+        Define a new marker link. This function must be called each time the number of markers change
+        Parameters
+        """
+
+        # Remove previous actors from the scene
+        for actor in self.markers_link_actors:
+            self.parent_window.ren.RemoveActor(actor)
+        self.markers_link_actors = list()
+
+        cells = vtkCellArray()
+        for i, j in enumerate(virtual_to_experimental_markers_indices):
+            if j is None:
+                continue
+            points = vtkPoints()
+            points.InsertNextPoint(self.markers["model"].data[:3, j, 0].data)
+            points.InsertNextPoint(self.markers["experimental"].data[:3, i, 0].data)
+
+            # Create the polygons
+            poly = vtkPolyLine()
+            poly.GetPointIds().SetNumberOfIds(2)  # make a line
+            poly.GetPointIds().SetId(0, 0)
+            poly.GetPointIds().SetId(1, 1)
+            cells.InsertNextCell(poly)
+
+            poly_data = vtkPolyData()
+            poly_data.SetPoints(points)
+            poly_data.SetLines(cells)
+
+            # Create a mapper
+            mapper = vtkPolyDataMapper()
+            mapper.SetInputData(poly_data)
+
+            # Create an actor
+            self.markers_link_actors.append(vtkActor())
+            self.markers_link_actors[-1].SetMapper(mapper)
+            self.markers_link_actors[-1].GetProperty().SetColor((1, 0, 0))
+
+            self.parent_window.ren.AddActor(self.markers_link_actors[-1])
+
+        # Update marker position
+        self._update_experimental_marker_link(virtual_to_experimental_markers_indices)
+
+    def _update_experimental_marker_link(self, virtual_to_experimental_markers_indices):
+        """
+        Update position of the line between experimental markers and virtual markers on the screen (but do not repaint)
+        Parameters
+        """
+
+        if len(self.markers_link_actors) != len(
+            tuple(i for i in virtual_to_experimental_markers_indices if i is not None)
+        ):
+            self._new_experimental_marker_link(virtual_to_experimental_markers_indices)
+            return  # Prevent calling update_markers recursively
+
+        cmp = 0
+        for i, j in enumerate(virtual_to_experimental_markers_indices):
+            if j is None:
+                continue
+            points = vtkPoints()
+            points.InsertNextPoint(self.markers["model"].data[:3, j, 0].data)
+            points.InsertNextPoint(self.markers["experimental"].data[:3, i, 0].data)
+
+            poly_line = self.markers_link_actors[cmp].GetMapper().GetInput()
+            poly_line.SetPoints(points)
+            cmp += 1
+
     def set_contacts_color(self, contacts_color):
         """
         Dynamically change the color of the contacts