Merge pull request #69 from pariterre/master

Easily manage the camera (command line) and take snapshot

bioviz/__init__.py

@@ -1,3 +1,4 @@
+from typing import Union
 import os
 import copy
 from functools import partial
@@ -316,6 +317,7 @@ def __init__(
         model_path=None,
         loaded_model=None,
         show_meshes=True,
+        mesh_opacity=0.8,
         show_global_center_of_mass=True,
         show_gravity_vector=True,
         show_floor=True,
@@ -387,6 +389,7 @@ def __init__(
             self.vtk_window,
             markers_color=(0, 0, 1),
             patch_color=InterfacesCollections.MeshColor.get_color(self.model),
+            mesh_opacity=mesh_opacity,
             markers_size=self.vtk_markers_size,
             contacts_size=contacts_size,
             segments_center_of_mass_size=segments_center_of_mass_size,
@@ -445,6 +448,7 @@ def __init__(
         # Create all the reference to the things to plot
         self.nQ = self.model.nbQ()
         self.Q = np.zeros(self.nQ)
+        self.idx_markers_to_remove = []
         if self.show_markers:
             self.Markers = InterfacesCollections.Markers(self.model)
             self.markers = Markers(np.ndarray((3, self.model.nbMarkers(), 1)))
@@ -463,6 +467,7 @@ def __init__(
             self.CoMbySegment = InterfacesCollections.CoMbySegment(self.model)
             self.segments_center_of_mass = Markers(np.ndarray((3, self.model.nbSegment(), 1)))
         if self.show_meshes:
+            self.show_segment_is_on = []
             self.mesh = []
             self.meshPointsInMatrix = InterfacesCollections.MeshPointsInMatrix(self.model)
             for i, vertices in enumerate(self.meshPointsInMatrix.get_data(Q=self.Q, compute_kin=False)):
@@ -472,6 +477,7 @@ def __init__(
                     else np.ndarray((0, 3), dtype="int32")
                 )
                 self.mesh.append(Mesh(vertex=vertices, triangles=triangles.T))
+                self.show_segment_is_on.append(True)
         if self.show_muscles:
             self.model.updateMuscles(self.Q, True)
             self.muscles = []
@@ -592,30 +598,18 @@ def set_q(self, Q, refresh_window=True):
         self.Q = Q
 
         self.model.UpdateKinematicsCustom(self.Q)
-        if self.show_muscles:
-            self.__set_muscles_from_q()
-        if self.show_local_ref_frame:
-            self.__set_rt_from_q()
-        if self.show_meshes:
-            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_floor:
-            self.__set_floor()
-        if self.show_segments_center_of_mass:
-            self.__set_segments_center_of_mass_from_q()
-        if self.show_markers:
-            self.__set_markers_from_q()
-        if self.show_markers:
-            self.__set_markers_from_q()
-        if self.show_contacts:
-            self.__set_contacts_from_q()
-        if self.show_soft_contacts:
-            self.__set_soft_contacts_from_q()
-        if self.show_wrappings:
-            self.__set_wrapping_from_q()
+
+        self.__set_muscles_from_q()
+        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()
+        self.__set_soft_contacts_from_q()
+        self.__set_wrapping_from_q()
 
         # Update the sliders
         if self.show_analyses_panel:
@@ -628,6 +622,53 @@ def set_q(self, Q, refresh_window=True):
         if refresh_window:
             self.refresh_window()
 
+    def get_camera_position(self) -> tuple:
+        return self.vtk_window.get_camera_position()
+
+    def set_camera_position(self, x: float, y: float, z: float):
+        self.vtk_window.set_camera_position(x, y, z)
+        self.refresh_window()
+
+    def get_camera_roll(self) -> float:
+        return self.vtk_window.get_camera_roll()
+
+    def set_camera_roll(self, roll: float):
+        self.vtk_window.set_camera_roll(roll)
+        self.refresh_window()
+
+    def get_camera_zoom(self) -> float:
+        return self.vtk_window.get_camera_zoom()
+
+    def set_camera_zoom(self, zoom: float):
+        self.vtk_window.set_camera_zoom(zoom)
+        self.refresh_window()
+
+    def get_camera_focus_point(self) -> tuple:
+        return self.vtk_window.get_camera_focus_point()
+
+    def set_camera_focus_point(self, x: float, y: float, z: float):
+        self.vtk_window.set_camera_focus_point(x, y, z)
+        self.refresh_window()
+
+    def toggle_segments(self, idx: Union[int, tuple]):
+        # Todo add graphical usage
+        if isinstance(idx, int):
+            idx = (idx,)
+        for i in idx:
+            self.show_segment_is_on[i] = not self.show_segment_is_on[i]
+        self.__set_meshes_from_q()
+
+        # Compute which marker index to remove
+        offset_marker = 0
+        self.idx_markers_to_remove = []
+        for s in range(self.model.nbSegment()):
+            nb_markers = self.model.nbMarkers(s)
+            if not self.show_segment_is_on[s]:
+                self.idx_markers_to_remove += list(range(offset_marker, offset_marker + nb_markers))
+            offset_marker += nb_markers
+        self.__set_markers_from_q()
+        self.__set_rt_from_q()
+
     def refresh_window(self):
         """
         Manually refresh the window. One should be aware when manually managing the window, that the plot won't even
@@ -931,11 +972,8 @@ def __animate_from_slider(self):
             self.Q = copy.copy(self.animated_Q[t, :])  # 1-based
             self.set_q(self.Q)
 
-        if self.show_experimental_markers:
-            self.__set_experimental_markers_from_frame()
-
-        if self.show_experimental_forces:
-            self.__set_experimental_forces_from_frame()
+        self.__set_experimental_markers_from_frame()
+        self.__set_experimental_forces_from_frame()
 
         # Update graph of muscle analyses
         self.__update_muscle_analyses_graphs(True, True, True, True)
@@ -962,6 +1000,15 @@ def __start_stop_animation(self):
             self.record_push_button.setEnabled(False)
             self.stop_record_push_button.setEnabled(False)
 
+    def snapshot(self, save_path: str):
+        # Todo Add a button
+        file_name, extension = os.path.splitext(save_path)
+        if not extension:
+            extension = ".png"
+        if extension != ".png":
+            raise NotImplementedError("The only snapshot format implemented is PNG")
+        self.vtk_window.snapshot(file_name + extension)
+
     def stop_recording(self):
         self._record(finish=True)
 
@@ -1131,15 +1178,26 @@ def load_experimental_forces(
             self.__start_stop_animation()
 
     def __set_markers_from_q(self):
+        if not self.show_markers:
+            return
+
         self.markers[0:3, :, :] = self.Markers.get_data(Q=self.Q, compute_kin=False)
+        if self.idx_markers_to_remove:
+            self.markers[0:3, self.idx_markers_to_remove, :] = np.nan
         self.vtk_model.update_markers(self.markers.isel(time=[0]))
 
     def __set_experimental_markers_from_frame(self):
+        if not self.show_experimental_markers:
+            return
+
         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]))
 
     def __set_experimental_forces_from_frame(self):
+        if not self.show_experimental_forces:
+            return
+
         segment_names = []
         for i in range(self.model.nbSegment()):
             segment_names.append(self.model.segment(i).name().to_string())
@@ -1176,19 +1234,31 @@ def __set_experimental_forces_from_frame(self):
         )
 
     def __set_contacts_from_q(self):
+        if not self.show_contacts:
+            return
+
         self.contacts[0:3, :, :] = self.Contacts.get_data(Q=self.Q, compute_kin=False)
         self.vtk_model.update_contacts(self.contacts.isel(time=[0]))
 
     def __set_soft_contacts_from_q(self):
+        if not self.show_soft_contacts:
+            return
+
         self.soft_contacts[0:3, :, :] = self.SoftContacts.get_data(Q=self.Q, compute_kin=False)
         self.vtk_model.update_soft_contacts(self.soft_contacts.isel(time=[0]))
 
     def __set_global_center_of_mass_from_q(self):
+        if not self.show_global_center_of_mass:
+            return
+
         com = self.CoM.get_data(Q=self.Q, compute_kin=False)
         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):
+        if not self.show_gravity_vector:
+            return
+
         start = [0, 0, 0]
         magnitude = self.Gravity.get_data()
         gravity = np.concatenate((start, magnitude))
@@ -1197,24 +1267,39 @@ def __set_gravity_vector(self):
         self.vtk_model.new_gravity_vector(id_matrix, gravity, length, normalization_ratio=0.3, vector_color=(0, 0, 0))
 
     def __set_floor(self):
+        if not self.show_floor:
+            return
+
         origin = self.floor_origin if self.floor_origin else (0, 0, 0)
         normal = self.floor_normal if self.floor_normal else self.Gravity.get_data()
         scale = self.floor_scale
         scale = (scale, scale, scale) if isinstance(scale, (int, float)) else scale
         self.vtk_model.new_floor(origin=origin, normal=normal, color=self.floor_color, scale=scale)
 
     def __set_segments_center_of_mass_from_q(self):
+        if not self.show_segments_center_of_mass:
+            return
+
         coms = self.CoMbySegment.get_data(Q=self.Q, compute_kin=False)
         for k, com in enumerate(coms):
             self.segments_center_of_mass.loc[{"channel": k, "time": 0}] = com.squeeze()
         self.vtk_model.update_segments_center_of_mass(self.segments_center_of_mass.isel(time=[0]))
 
     def __set_meshes_from_q(self):
+        if not self.show_meshes:
+            return
+
         for m, meshes in enumerate(self.meshPointsInMatrix.get_data(Q=self.Q, compute_kin=False)):
-            self.mesh[m][0:3, :, :] = meshes
+            if self.show_segment_is_on[m]:
+                self.mesh[m][0:3, :, :] = meshes
+            else:
+                self.mesh[m][0:3, :, :] = np.nan
         self.vtk_model.update_mesh(self.mesh)
 
     def __set_muscles_from_q(self):
+        if not self.show_muscles:
+            return
+
         muscles = self.musclesPointsInGlobal.get_data(Q=self.Q)
         idx = 0
         cmp = 0
@@ -1228,6 +1313,9 @@ def __set_muscles_from_q(self):
         self.vtk_model.update_muscle(self.muscles)
 
     def __set_wrapping_from_q(self):
+        if not self.show_wrappings:
+            return
+
         for i, wraps in enumerate(self.wraps_base):
             for j, wrap in enumerate(wraps):
                 if self.model.muscle(i).pathModifier().object(j).typeOfNode() == biorbd.WRAPPING_HALF_CYLINDER:
@@ -1242,6 +1330,12 @@ def __set_wrapping_from_q(self):
         self.vtk_model.update_wrapping(self.wraps_current)
 
     def __set_rt_from_q(self):
+        if not self.show_local_ref_frame:
+            return
+
         for k, rt in enumerate(self.allGlobalJCS.get_data(Q=self.Q, compute_kin=False)):
-            self.rt[k] = Rototrans(rt)
+            if self.show_segment_is_on[k]:
+                self.rt[k] = Rototrans(rt)
+            else:
+                self.rt[k] = Rototrans(np.eye(4)) * np.nan
         self.vtk_model.update_rt(self.rt)

bioviz/_version.py

@@ -1 +1 @@
-__version__ = "2.1.6"
+__version__ = "2.1.7"

bioviz/biorbd_vtk.py

@@ -1,7 +1,7 @@
 """
 Visualization toolkit in pyomeca
 """
-
+import os
 import time
 import sys
 
@@ -33,6 +33,7 @@
     vtkTransform,
     vtkTransformPolyDataFilter,
     vtkPlaneSource,
+    vtkPNGWriter,
 )
 
 from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
@@ -109,6 +110,35 @@ def update_frame(self):
         self.interactor.Render()
         app.processEvents()
 
+    def get_camera_position(self) -> tuple:
+        return self.ren.GetActiveCamera().GetPosition()
+
+    def set_camera_position(self, x: float, y: float, z: float):
+        cam = self.ren.GetActiveCamera()
+        cam.SetPosition(x, y, z)
+        self.ren.ResetCamera()
+
+    def get_camera_roll(self) -> float:
+        return self.ren.GetActiveCamera().GetRoll()
+
+    def set_camera_roll(self, roll: float):
+        cam = self.ren.GetActiveCamera()
+        cam.SetRoll(roll)
+        self.ren.ResetCamera()
+
+    def get_camera_zoom(self) -> float:
+        return 1 / self.ren.GetActiveCamera().GetParallelScale()
+
+    def set_camera_zoom(self, zoom: float):
+        cam = self.ren.GetActiveCamera()
+        cam.SetParallelScale(1 / zoom)
+
+    def get_camera_focus_point(self) -> tuple:
+        return self.ren.GetActiveCamera().GetFocalPoint()
+
+    def set_camera_focus_point(self, x: float, y: float, z: float):
+        self.ren.GetActiveCamera().SetFocalPoint(x, y, z)
+
     def change_background_color(self, color):
         """
         Dynamically change the background color of the windows
@@ -119,6 +149,23 @@ def change_background_color(self, color):
         self.ren.SetBackground(color)
         self.setPalette(QPalette(QColor(int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))))
 
+    def snapshot(self, save_path: str):
+        w2if = vtkWindowToImageFilter()
+        w2if.SetInput(self.avatar_widget.GetRenderWindow())
+        w2if.Update()
+
+        w = vtkPNGWriter()
+        folder_path = os.path.dirname(save_path)
+        try:
+            os.mkdir(folder_path)
+        except FileExistsError:
+            pass
+
+        w.SetFileName(save_path)
+        w.SetInputData(w2if.GetOutput())
+
+        w.Write()
+
     def record(self, finish=False, button_to_block=(), file_name=None):
         windowToImageFilter = vtkWindowToImageFilter()
         self.video_recorder.SetInputConnection(windowToImageFilter.GetOutputPort())