diff --git a/bioviz/__init__.py b/bioviz/__init__.py index 33f942d..7be658f 100644 --- a/bioviz/__init__.py +++ b/bioviz/__init__.py @@ -301,6 +301,7 @@ def __init__( loaded_model=None, show_meshes=True, show_global_center_of_mass=True, + show_gravity_vector=True, show_segments_center_of_mass=True, segments_center_of_mass_size=0.005, show_global_ref_frame=True, @@ -399,6 +400,7 @@ def __init__( self.soft_contacts_color = soft_contacts_color self.show_global_ref_frame = show_global_ref_frame self.show_global_center_of_mass = show_global_center_of_mass + self.show_gravity_vector = show_gravity_vector self.show_segments_center_of_mass = show_segments_center_of_mass self.show_local_ref_frame = show_local_ref_frame self.biorbd_compiled_with_muscles = hasattr(biorbd.Model, "nbMuscles") @@ -572,6 +574,8 @@ def set_q(self, Q, refresh_window=True): self.__set_meshes_from_q() if self.show_global_center_of_mass: self.__set_global_center_of_mass_from_q() + if self.show_gravity_vector: + self.__set_gravity_vector() if self.show_segments_center_of_mass: self.__set_segments_center_of_mass_from_q() if self.show_markers: @@ -1135,6 +1139,16 @@ def __set_global_center_of_mass_from_q(self): self.global_center_of_mass.loc[{"channel": 0, "time": 0}] = com.squeeze() self.vtk_model.update_global_center_of_mass(self.global_center_of_mass.isel(time=[0])) + def __set_gravity_vector(self): + start = [0, 0, 0] + magnitude = self.model.getGravity().to_array() + gravity = np.concatenate((start, magnitude)) + length = np.linalg.norm(gravity) + id_matrix = np.identity(4) + self.vtk_model.new_gravity_vector( + id_matrix, gravity, length, normalization_ratio=0.3, vector_color=(0, 0, 0) + ) + def __set_segments_center_of_mass_from_q(self): coms = self.CoMbySegment.get_data(Q=self.Q, compute_kin=False) for k, com in enumerate(coms): @@ -1148,7 +1162,6 @@ def __set_meshes_from_q(self): def __set_muscles_from_q(self): muscles = self.musclesPointsInGlobal.get_data(Q=self.Q) - idx = 0 cmp = 0 for group_idx in range(self.model.nbMuscleGroups()): diff --git a/bioviz/biorbd_vtk.py b/bioviz/biorbd_vtk.py index c4da4f9..4e80151 100644 --- a/bioviz/biorbd_vtk.py +++ b/bioviz/biorbd_vtk.py @@ -271,6 +271,8 @@ def __init__( self.normalization_ratio = 0.2 self.force_actors = list() + self.gravity_actors = list() + def set_markers_color(self, markers_color): """ Dynamically change the color of the markers @@ -1346,7 +1348,6 @@ def set_force_color(self, force_color): 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): """ @@ -1360,7 +1361,6 @@ def set_force_opacity(self, force_opacity): """ 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): """ @@ -1514,3 +1514,70 @@ def update_force(self, segment_jcs, all_forces, max_forces, normalization_ratio) self.force_actors[i].SetMapper(mapper) self.force_actors[i].GetProperty().SetColor(self.force_color) self.force_actors[i].GetProperty().SetOpacity(self.force_opacity) + + def new_gravity_vector(self, segment_rt, gravity, length, normalization_ratio, vector_color): + """ + Define a new gravity vector. + Parameters + ---------- + segment_rt : np.ndarray + homogeneous matrix in which coordinates are applied + gravity : np.ndarray + gravity array with 3 application coordinates and 3 magnitude coordinates + length : float + absolut length of gravity vector to scale arrow + normalization_ratio: float + ratio to scale arrow + vector_color: tuple + tuple of RGB for vector color + + """ + # 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 + + # Arrow orientation + transform = vtkTransform() + transform_polydata = vtkTransformPolyDataFilter() + transform_polydata.SetTransform(transform) + transform_polydata.SetInputConnection(arrow_source.GetOutputPort()) + + rot_seg = segment_rt[:3, :3] + trans_seg = segment_rt[:-1, 3:] + force_magnitude = np.dot(rot_seg, gravity[3:]) + force_magnitude = force_magnitude + trans_seg.reshape(3) + force_application = np.dot(rot_seg, gravity[:3]) + force_application = force_application + trans_seg.reshape(3) + 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 / length + transform = vtkTransform() + transform.Translate(application_point) + transform.Concatenate(matrix) + transform.Scale(length, length, length) + + # Create an actor + mapper = vtkPolyDataMapper() + + transform_polydata = vtkTransformPolyDataFilter() + transform_polydata.SetTransform(transform) + transform_polydata.SetInputConnection(self.arrow_source.GetOutputPort()) + mapper.SetInputConnection(transform_polydata.GetOutputPort()) + + self.gravity_actors = vtkActor() + self.gravity_actors.SetMapper(mapper) + self.gravity_actors.GetProperty().SetColor(vector_color) + + self.parent_window.ren.AddActor(self.gravity_actors) + self.parent_window.ren.ResetCamera()