diff --git a/bioviz/__init__.py b/bioviz/__init__.py index cf6fca6..59f679e 100644 --- a/bioviz/__init__.py +++ b/bioviz/__init__.py @@ -7,6 +7,7 @@ import scipy import xarray as xr import pandas + try: import biorbd except ImportError: @@ -280,6 +281,7 @@ def __init__( show_analyses_panel=True, background_color=(0.5, 0.5, 0.5), force_wireframe=False, + experimental_forces_colors=(85, 78, 0), **kwargs, ): """ @@ -321,6 +323,7 @@ def __init__( contacts_size=contacts_size, segments_center_of_mass_size=segments_center_of_mass_size, force_wireframe=force_wireframe, + force_color=experimental_forces_colors, ) self.vtk_model_markers: VtkModel = None self.is_executing = False @@ -337,6 +340,11 @@ def __init__( self.show_experimental_markers = False self.experimental_markers = None self.experimental_markers_color = experimental_markers_color + self.show_experimental_forces = False + self.experimental_forces = None + self.segment_forces = [] + self.experimental_forces_color = experimental_forces_colors + self.force_normalization_ratio = None self.show_contacts = show_contacts self.show_global_ref_frame = show_global_ref_frame @@ -828,6 +836,9 @@ def __animate_from_slider(self): if self.show_experimental_markers: self.__set_experimental_markers_from_frame() + if self.show_experimental_forces: + self.__set_experimental_forces_from_frame() + # Update graph of muscle analyses self.__update_muscle_analyses_graphs(True, True, True, True) @@ -923,7 +934,12 @@ def __set_movement_slider(self): # Update the slider bar and frame count self.movement_slider[0].setEnabled(True) self.movement_slider[0].setMinimum(1) - experiment_shape = 0 if self.experimental_markers is None else self.experimental_markers.shape[2] + experiment_shape = 0 + if self.experimental_forces is not None: + experiment_shape = self.experimental_forces.shape[2] + if self.experimental_markers is not None: + experiment_shape = max(self.experimental_markers.shape[2], experiment_shape) + q_shape = 0 if self.animated_Q is None else self.animated_Q.shape[0] self.movement_slider[0].setMaximum(max(q_shape, experiment_shape)) pal = QPalette() @@ -937,9 +953,7 @@ def __load_experimental_data_from_button(self): # Load the actual movement options = QFileDialog.Options() options |= QFileDialog.DontUseNativeDialog - file_name = QFileDialog.getOpenFileName( - self.vtk_window, "Data to load", "", "C3D (*.c3d)", options=options - ) + file_name = QFileDialog.getOpenFileName(self.vtk_window, "Data to load", "", "C3D (*.c3d)", options=options) if not file_name[0]: return self.load_experimental_markers(file_name[0]) @@ -958,7 +972,7 @@ def load_experimental_markers(self, data, auto_start=True, ignore_animation_warn else: raise RuntimeError( - f"Wrong type of experimental markers data ({type(data)}. " + f"Wrong type of experimental markers data ({type(data)}). " f"Allowed type are numpy array (3xNxT), data array (3xNxT) or .c3d file (str)." ) @@ -974,6 +988,38 @@ def load_experimental_markers(self, data, auto_start=True, ignore_animation_warn if auto_start: self.__start_stop_animation() + def load_experimental_forces( + self, data, segments=None, normalization_ratio=0.2, auto_start=True, ignore_animation_warning=True + ): + if isinstance(data, (np.ndarray, xr.DataArray)): + self.experimental_forces = data if isinstance(data, xr.DataArray) else xr.DataArray(data) + else: + raise RuntimeError( + f"Wrong type of experimental force data ({type(data)}). " + f"Allowed type are numpy array (SxNxT), data array (SxNxT)." + ) + + if segments: + self.segment_forces = segments if isinstance(segments, (list, np.ndarray)) else [segments] + else: + self.segment_forces = ["ground"] * self.experimental_forces.shape[0] + + if len(self.segment_forces) != self.experimental_forces.shape[0]: + raise RuntimeError( + "Number of segment must match number of experimental forces. " + f"You have {len(segments)} and {self.experimental_forces.shape[0]}." + ) + + self.force_normalization_ratio = normalization_ratio + + self.show_experimental_forces = True + self.__set_movement_slider() + + if ignore_animation_warning: + self.animation_warning_already_shown = True + if auto_start: + self.__start_stop_animation() + def __set_markers_from_q(self): self.markers[0:3, :, :] = self.Markers.get_data(Q=self.Q, compute_kin=False) self.vtk_model.update_markers(self.markers.isel(time=[0])) @@ -981,7 +1027,43 @@ def __set_markers_from_q(self): 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.isel(time=t)) + self.vtk_model_markers.update_markers(self.experimental_markers[:, :, t : t + 1].isel(time=[0])) + + def __set_experimental_forces_from_frame(self): + segment_names = [] + for i in range(self.model.nbSegment()): + segment_names.append(self.model.segment(i).name().to_string()) + global_jcs = self.allGlobalJCS.get_data(Q=self.Q, compute_kin=False) + segment_jcs = [] + + for segment in self.segment_forces: + if isinstance(segment, str): + if segment == "ground": + segment_jcs.append(np.identity(4)) + else: + segment_jcs.append(global_jcs[segment_names.index(segment)]) + elif isinstance(segment, (float, int)): + segment_jcs.append(global_jcs[segment]) + else: + raise RuntimeError("Wrong type of segment.") + + max_forces = [] + for i, forces in enumerate(self.experimental_forces): + max_forces.append( + max( + np.sqrt( + (forces[3, :] - forces[0, :]) ** 2 + + (forces[4, :] - forces[1, :]) ** 2 + + (forces[5, :] - forces[2, :]) ** 2 + ) + ) + ) + + t_slider = self.movement_slider[0].value() - 1 + t = t_slider if t_slider < self.experimental_forces.shape[2] else self.experimental_forces.shape[2] - 1 + self.vtk_model.update_force( + segment_jcs, self.experimental_forces[:, :, t : t + 1], max_forces, self.force_normalization_ratio + ) def __set_contacts_from_q(self): self.contacts[0:3, :, :] = self.Contacts.get_data(Q=self.Q, compute_kin=False) diff --git a/bioviz/biorbd_vtk.py b/bioviz/biorbd_vtk.py index e379a2b..05df66f 100644 --- a/bioviz/biorbd_vtk.py +++ b/bioviz/biorbd_vtk.py @@ -25,7 +25,15 @@ vtkWindowToImageFilter, vtkPolygon, vtkExtractEdges, + vtkArrowSource, + vtkNamedColors, + vtkMath, + vtkMatrix4x4, + vtkMinimalStandardRandomSequence, + vtkTransform, + vtkTransformPolyDataFilter, ) + from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor from pyomeca import Markers, Rototrans @@ -168,6 +176,8 @@ def __init__( muscle_opacity=1.0, rt_length=0.1, rt_width=2, + force_color=(85, 78, 0), + force_opacity=1.0, ): """ Creates a model that will holds things to plot @@ -244,6 +254,14 @@ def __init__( self.wrapping_opacity = wrapping_opacity self.wrapping_actors = list() + self.all_forces = [] + self.force_centers = [] + self.max_forces = [] + self.force_color = force_color + self.force_opacity = force_opacity + self.normalization_ratio = 0.2 + self.force_actors = list() + def set_markers_color(self, markers_color): """ Dynamically change the color of the markers @@ -1212,3 +1230,181 @@ def create_global_ref_frame(self): self.parent_window.ren.AddActor(actor) self.parent_window.ren.ResetCamera() + + def set_force_color(self, force_color): + """ + Dynamically change the color of the force + Parameters + ---------- + force_color : tuple(int) + Color the force should be drawn (1 is max brightness) + """ + self.force_color = force_color + self.update_force(self.force_centers, self.all_force, self.max_forces, self.normalization_ratio) + + def set_force_opacity(self, force_opacity): + """ + Dynamically change the opacity of the _force + Parameters + ---------- + force_opacity : float + Opacity of the force (0.0 is completely transparent, 1.0 completely opaque) + Returns + ------- + + """ + self.force_opacity = force_opacity + self.update_force(self.force_centers, self.all_force, self.max_forces, self.normalization_ratio) + + def new_force_set(self, segment_jcs, all_forces, max_forces, normalization_ratio): + """ + Define a new force set. + Parameters + ---------- + segment_jcs : list + list of transformation for each segment + all_forces : np.ndarray + force array. Dims are : (1) segments, (2) application and magnitude point arrow coordinates (3) frame + max_forces : list + list of maximal force for each segment on all frames + + """ + # Arrow visualization parameters + arrow_source = vtkArrowSource() + arrow_source.SetTipResolution(15) + arrow_source.SetShaftResolution(8) + arrow_source.SetShaftRadius(0.015) + arrow_source.SetTipLength(0.2) + arrow_source.SetTipRadius(0.08) + + self.arrow_source = arrow_source + + transform = vtkTransform() + + transform_polydata = vtkTransformPolyDataFilter() + transform_polydata.SetTransform(transform) + transform_polydata.SetInputConnection(arrow_source.GetOutputPort()) + + self.all_forces = all_forces + # Remove previous actors from the scene + for actor in self.force_actors: + self.parent_window.ren.RemoveActor(actor) + self.force_actors = list() + + for i, forces in enumerate(all_forces): + # Create a mapper + mapper = vtkPolyDataMapper() + mapper.SetInputConnection(transform_polydata.GetOutputPort()) + + # Create an actor + self.force_actors.append(vtkActor()) + self.force_actors[i].SetMapper(mapper) + self.force_actors[i].GetProperty().SetColor(self.force_color) + self.force_actors[i].GetProperty().SetOpacity(self.force_opacity) + + self.parent_window.ren.AddActor(self.force_actors[i]) + self.parent_window.ren.ResetCamera() + + # Set rt orientations + self.n_force = len(all_forces) + self.update_force(segment_jcs, all_forces, max_forces, normalization_ratio) + + @staticmethod + def compute_basis_force(application_point, magnitude_point): + """ + Compute basis to plot vtk arrow object from two points. + Parameters + ---------- + application_point : list + list of the 3 coordinates for the arrow starting point + magnitude_point: list + list of the 3 coordinates for the arrow ending point + Return + ---------- + matrix of the transformation and the length of the arrow. + """ + # Compute a basis + normalizedX = [0.0] * 3 + normalizedY = [0.0] * 3 + normalizedZ = [0.0] * 3 + + # The X axis is a vector from start to end + vtkMath.Subtract(magnitude_point, application_point, normalizedX) + length = vtkMath.Norm(normalizedX) + vtkMath.Normalize(normalizedX) + + rng = vtkMinimalStandardRandomSequence() + rng.SetSeed(8775070) + max_r = 10.0 + + # The Z axis is an arbitrary vector cross X + arbitrary = [0.0] * 3 + for i in range(0, 3): + arbitrary[i] = rng.GetRangeValue(-max_r, max_r) + rng.Next() + vtkMath.Cross(normalizedX, arbitrary, normalizedZ) + vtkMath.Normalize(normalizedZ) + + # The Y axis is Z cross X + vtkMath.Cross(normalizedZ, normalizedX, normalizedY) + matrix = vtkMatrix4x4() + + # Create the direction cosine matrix + matrix.Identity() + for i in range(0, 3): + matrix.SetElement(i, 0, normalizedX[i]) + matrix.SetElement(i, 1, normalizedY[i]) + matrix.SetElement(i, 2, normalizedZ[i]) + + return matrix, length + + def update_force(self, segment_jcs, all_forces, max_forces, normalization_ratio): + """ + Update force on the screen (but do not repaint) + Parameters + ---------- + segment_jcs : list + list of roto-translation matrix for each segment + all_forces : np.ndarray + force array. Dims are : (1) nb of segments, (2) arrow start and end points coordinates (3) frame + max_forces : list + list of maximal force for each segment compute from all frames + normalization_ratio : float + ratio to normalize force for visualization + """ + if self.all_forces == []: + self.new_force_set(segment_jcs, all_forces, max_forces, normalization_ratio) + return # Prevent calling update_markers recursively + + self.max_forces = max_forces + self.all_forces = all_forces + for i, forces in enumerate(all_forces): + # Express force from current segment basis to global basis + rot_seg = segment_jcs[i][:3, :3] + trans_seg = segment_jcs[i][:-1, 3:] + force_magnitude = np.dot(rot_seg, forces[3:, :]) + force_magnitude = force_magnitude + trans_seg + force_application = np.dot(rot_seg, forces[:3, :]) + force_application = force_application + trans_seg + application_point = [force_application[0, :], force_application[1, :], force_application[2, :]] + magnitude_point = [force_magnitude[0, :], force_magnitude[1, :], force_magnitude[2, :]] + # Compute a basis for the arrow scaling + matrix, length = self.compute_basis_force(application_point, magnitude_point) + # Normalize force for visualization + length = length * normalization_ratio / max_forces[i] + transform = vtkTransform() + transform.Translate(application_point) + transform.Concatenate(matrix) + transform.Scale(length, length, length) + + # Create an actor + mapper = self.force_actors[i].GetMapper() + transform_polydata = vtkTransformPolyDataFilter() + transform_polydata.SetTransform(transform) + transform_polydata.SetInputConnection(self.arrow_source.GetOutputPort()) + + mapper.SetInputConnection(transform_polydata.GetOutputPort()) + + self.force_actors[i].SetMapper(mapper) + self.force_actors[i].GetProperty().SetColor(self.force_color) + self.force_actors[i].GetProperty().SetOpacity(self.force_opacity) diff --git a/bioviz/mesh.py b/bioviz/mesh.py index 0edcadc..c4540ca 100644 --- a/bioviz/mesh.py +++ b/bioviz/mesh.py @@ -3,12 +3,7 @@ class Mesh(Markers): - def __new__( - cls, - vertex=None, - triangles=np.ndarray((3, 0)), - **kwargs, - ): + def __new__(cls, vertex=None, triangles=np.ndarray((3, 0)), **kwargs): if isinstance(triangles, list): triangles = np.array(triangles)