diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 019421b3..b37cd3ad 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -14,5 +14,5 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v4 with: - python-version: 3.7 + python-version: 3.8 - uses: isort/isort-action@master diff --git a/README.md b/README.md index 71d2b8a8..565320a6 100644 --- a/README.md +++ b/README.md @@ -1,76 +1,107 @@ -[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) -[![Tests](https://github.com/eth-ait/aitviewer/actions/workflows/tests.yml/badge.svg)](https://github.com/eth-ait/aitviewer/actions/workflows/tests.yml) -[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.10013305.svg)](https://doi.org/10.5281/zenodo.10013305) +# aitviewer - SKEL + +This fork of AitViewer enables the vizualization of Marker sequences, OpenSim models sequences, the BSM model and the SKEL model. + +This repo contain a visualization tool. If you are interested in the SKEL model code, please refer to the [SKEL repository](https://download.is.tue.mpg.de/skel/main_paper.pdf). + +For more info on SKEL, BSM and BioaAmass, check our [project page](https://skel.is.tue.mpg.de) and our [paper](https://download.is.tue.mpg.de/skel/main_paper.pdf). + +aitviewer is a set of tools to visualize and interact with sequences of 3D data with cross-platform support on Windows, Linux, and macOS. See the official page at [https://eth-ait.github.io/aitviewer](https://eth-ait.github.io/aitviewer/) for all the details. + +![aitviewer skel gif](assets/skel_sequence.gif) +⇧ *aitviewer-Skel enables visualization of motion sequences of the SKEL model.* -# [![aitviewer](assets/aitviewer_logo.svg)](https://github.com/eth-ait/aitviewer) -A set of tools to visualize and interact with sequences of 3D data with cross-platform support on Windows, Linux, and macOS. See the official page at [https://eth-ait.github.io/aitviewer](https://eth-ait.github.io/aitviewer/) for all the details. ## Installation -Basic Installation: -```commandline -pip install aitviewer -``` -Note that this does not install the GPU-version of PyTorch automatically. If your environment already contains it, you should be good to go, otherwise install it manually. -Or install locally (if you need to extend or modify code) -```commandline -git clone git@github.com:eth-ait/aitviewer.git -cd aitviewer +Clone this repository and install it using: +``` +git clone https://github.com/MarilynKeller/aitviewer-skel.git +cd aitviewer-skel pip install -e . ``` -On macOS with Apple Silicon it is recommended to use PyQt6. Please check [this issue](https://github.com/eth-ait/aitviewer/issues/22) for installation instructions. +To set up the paths to SMPLX and AMASS, please refer to the [aitviewer instructions](https://eth-ait.github.io/aitviewer/frontend.html#configure-the-viewer) + +## BSM model + +You can download the BSM model `bsm.osim` from the dowload page at [https://skel.is.tue.mpg.de](https://skel.is.tue.mpg.de). To visualize it, run: + +```python load_osim.py --osim /path/to/bsm.osim``` + +You can find motion sequences in the BioAmass dataset at [https://skel.is.tue.mpg.de](https://skel.is.tue.mpg.de). + +To visualize an OpenSim motion sequence: + +``` +python load_osim.py --osim /path/to/bsm.osim --mot /path/to/trial.mot +``` + +![aitviewer osim vizu](assets/osim_apose.png) + +## SKEL model + +You can download the SKEL model from the dowload page at [https://skel.is.tue.mpg.de](https://skel.is.tue.mpg.de). +Edit then the file aitviewer/aitvconfig.yaml` to point to the SKEL folder: +```skel_models: "/path/to/skel_models_v1.0"``` + +Install the SKEL loader by following the guidelines here: https://github.com/MarilynKeller/SKEL + +Vizualize the SKEL model's shape space: + +``` +python examples/load_SKEL.py +``` + +Vizualize a SKEL sequence. You can find a sample SKEL motion in `skel_models_v1.0/sample_motion/ ` and the corresponding SMPL motion. -For more advanced installation and for installing SMPL body models, please refer to the [documentation](https://eth-ait.github.io/aitviewer/parametric_human_models/supported_models.html) . +``` +python examples/load_SKEL.py -s 'skel_models_v1.1/sample_motion/01_01_poses_skel.pkl' --z_up +``` -## Features -* Native Python interface, easy to use and hack. -* Load [SMPL[-H/-X]](https://smpl.is.tue.mpg.de/) / [MANO](https://mano.is.tue.mpg.de/) / [FLAME](https://flame.is.tue.mpg.de/) / [STAR](https://github.com/ahmedosman/STAR) / [SUPR](https://github.com/ahmedosman/SUPR) sequences and display them in an interactive viewer. -* Headless mode for server rendering of videos/images. -* Remote mode for non-blocking integration of visualization code. -* Render 3D data on top of images via weak-perspective or OpenCV camera models. -* Animatable camera paths. -* Edit SMPL sequences and poses manually. -* Prebuilt renderable primitives (cylinders, spheres, point clouds, etc). -* Built-in extensible GUI (based on Dear ImGui). -* Export screenshots, videos and turntable views (as mp4/gif) -* High-Performance ModernGL-based rendering pipeline (running at 100fps+ on most laptops). -![aitviewer SMPL Editing](https://user-images.githubusercontent.com/5639197/188625764-351100e9-992e-430c-b170-69d4f142f5dd.gif) +## BioAmass Dataset +First download the models and dataset from [https://skel.is.tue.mpg.de](https://skel.is.tue.mpg.de) and in `aitconfig.yaml` set the following paths: +``` +osim_geometry : /path/to/skel_models_v1.0/Geometry +bioamass : /path/to/bioamass_v1.0 +``` -## Quickstart -Display an SMPL T-pose (Requires SMPL models): -```py -from aitviewer.renderables.smpl import SMPLSequence -from aitviewer.viewer import Viewer +To visualize a sequence from the BioAmass dataset, run: -if __name__ == '__main__': - v = Viewer() - v.scene.add(SMPLSequence.t_pose()) - v.run() ``` +python examples/load_bioamass.py +``` + +## Mocap data + +We enable loading .c3d and .trc motion capture data. Sample CMU mocap data can be downloaded at http://mocap.cs.cmu.edu/subjects.php. Set the path to the mocap data folder in `aitvconfig.yaml` in `datasets.mocap`. +To visualize an example mocap sequence, run: -## Projects using the aitviewer -A sampling of projects using the aitviewer. Let us know if you want to be added to this list! -- Fan et al., [HOLD: Category-agnostic 3D Reconstruction of Interacting Hands and Objects from Video](https://github.com/zc-alexfan/hold), CVPR 2024 -- Braun et al., [Physically Plausible Full-Body Hand-Object Interaction Synthesis](https://eth-ait.github.io/phys-fullbody-grasp/), 3DV 2024 -- Zhang and Christen et al., [ArtiGrasp: Physically Plausible Synthesis of Bi-Manual Dexterous Grasping and Articulation](https://eth-ait.github.io/artigrasp/), 3DV 2024 -- Kaufmann et al., [EMDB: The Electromagnetic Database of Global 3D Human Pose and Shape in the Wild](https://ait.ethz.ch/emdb), ICCV 2023 -- Shen and Guo et al., [X-Avatar: Expressive Human Avatars](https://skype-line.github.io/projects/X-Avatar/), CVPR 2023 -- Sun et al., [TRACE: 5D Temporal Regression of Avatars with Dynamic Cameras in 3D Environments](https://www.yusun.work/TRACE/TRACE.html), CVPR 2023 -- Fan et al., [ARCTIC: A Dataset for Dexterous Bimanual Hand-Object Manipulation](https://github.com/zc-alexfan/arctic), CVPR 2023 -- Dong and Guo et al., [PINA: Learning a Personalized Implicit Neural Avatar from a Single RGB-D Video Sequence](https://zj-dong.github.io/pina/), CVPR 2022 -- Dong et al., [Shape-aware Multi-Person Pose Estimation from Multi-view Images](https://ait.ethz.ch/projects/2021/multi-human-pose/), ICCV 2021 -- Kaufmann et al., [EM-POSE: 3D Human Pose Estimation from Sparse Electromagnetic Trackers](https://ait.ethz.ch/projects/2021/em-pose/), ICCV 2021 -- Vechev et al., [Computational Design of Kinesthetic Garments](https://ait.ethz.ch/projects/2022/cdkg/), Eurographics 2021 -- Guo et al., [Human Performance Capture from Monocular Video in the Wild](https://ait.ethz.ch/projects/2021/human-performance-capture/index.php), 3DV 2021 +```python load_markers.py``` + + ## Citation -If you use this software, please cite it as below. -```commandline +If you use this software, please cite the following work and software: + +``` +@inproceedings{keller2023skel, + title = {From Skin to Skeleton: Towards Biomechanically Accurate 3D Digital Humans}, + author = {Keller, Marilyn and Werling, Keenon and Shin, Soyong and Delp, Scott and + Pujades, Sergi and Liu, C. Karen and Black, Michael J.}, + booktitle = {ACM ToG, Proc.~SIGGRAPH Asia}, + volume = {42}, + number = {6}, + month = dec, + year = {2023}, +} +``` + +``` @software{Kaufmann_Vechev_aitviewer_2022, author = {Kaufmann, Manuel and Vechev, Velko and Mylonopoulos, Dario}, doi = {10.5281/zenodo.10013305}, @@ -81,9 +112,12 @@ If you use this software, please cite it as below. } ``` -## Contact & Contributions -This software was developed by [Manuel Kaufmann](mailto:manuel.kaufmann@inf.ethz.ch), [Velko Vechev](mailto:velko.vechev@inf.ethz.ch) and Dario Mylonopoulos. -For questions please create an issue. -We welcome and encourage module and feature contributions from the community. +## Licencing + +For use of SKEL and BSM, please refer to our project page https://skel.is.tue.mpg.de/license.html. + +## Contact + +For any question on the OpenSim model or SKEL loading, please contact skel@tuebingen.mpg.de. -![aitviewer Sample](assets/aitviewer_collab.png) +For commercial licensing, please contact ps-licensing@tue.mpg.de diff --git a/aitviewer/aitvconfig.yaml b/aitviewer/aitvconfig.yaml index 283f25d2..402402ea 100644 --- a/aitviewer/aitvconfig.yaml +++ b/aitviewer/aitvconfig.yaml @@ -1,12 +1,18 @@ # Access SMPL models. -smplx_models: "../data/smplx_models" +smplx_models: "/is/cluster/fast/mkeller2/Data/body_model/smplx_models" #"/ps/project/rib_cage_breathing/TML/Data/smplx_models" star_models: "../data/star_models" supr_models: "../data/supr_models" +skel_models: "/is/cluster/work/mkeller2/Data/TML/Release/v1.0/skel_models_v1.1" +osim_geometry : "/is/cluster/work/mkeller2/Data/TML/Release/v1.0/skel_models_v1.1/Geometry" # Access to datasets. datasets: amass: - "../data/amass" + "/is/cluster/work/mkeller2/Data/TML/AMASS" + bioamass: + "/is/cluster/work/mkeller2/Data/TML/Release/v1.0/bioamass_v1.0" + amass_mocap: + "/is/cluster/work/mkeller2/Data/TML/AMASS_mocap" # Pytorch configuration. device: "cuda:0" @@ -34,11 +40,12 @@ auto_set_floor: True auto_set_camera_target: True backface_culling: True background_color: [1.0, 1.0, 1.0, 1.0] -window_type: "pyqt5" +window_type: "pyglet" # window_type: "pyqt5" + # Viewer defaults to Y up, set to true for Z up. z_up: False # Server for remote connections. -server_enabled: False +server_enabled: True server_port: 8417 diff --git a/aitviewer/renderables/markers.py b/aitviewer/renderables/markers.py new file mode 100644 index 00000000..602c45e8 --- /dev/null +++ b/aitviewer/renderables/markers.py @@ -0,0 +1,239 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import os +import pickle as pkl + +import numpy as np +import tqdm + +from aitviewer.renderables.point_clouds import PointClouds +from aitviewer.renderables.spheres import Spheres +from aitviewer.scene.node import Node +from aitviewer.utils.mocap import load_markers +from aitviewer.configuration import CONFIG as C + +class Markers(Node): + """ + Draw a point clouds man! + """ + + def __init__( + self, + points, + markers_labels, + name="Mocap data", + colors=None, + lengths=None, + point_size=5.0, + radius=0.0075, + color=(0.0, 0.0, 1.0, 1.0), + as_spheres=True, + z_up = False, + **kwargs, + ): + """ + A sequence of point clouds. Each point cloud can have a varying number of points. + Internally represented as a list of arrays. + :param points: Sequence of points (F, P, 3) + :param colors: Sequence of Colors (F, C, 4) + :param lengths: Length mask for each frame of points denoting the usable part of the array + :param point_size: Initial point size + """ + # self.points = points + super(Markers, self).__init__(name, n_frames=points.shape[0], color=color, **kwargs) + + # Check that the marker labels are sorted + # markers_labels_copy = markers_labels.copy() + # markers_labels_copy.sort() + # assert markers_labels == markers_labels_copy + + self.markers_labels = markers_labels + self.marker_trajectory = points # FxMx3 + self.color = color + + self._z_up = z_up + + if self._z_up and not C.z_up: + self.rotation = np.matmul(np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]), self.rotation) + + + if self.marker_trajectory.shape[1] > 200: + as_spheres = False + print(f"Too many markers ({self.marker_trajectory.shape[1]}). Switching to pointcloud.") + + # todo fix color bug + for mi, marker_name in enumerate(self.markers_labels): + if colors is not None: + color = tuple(colors[mi]) + + if as_spheres: + markers_seq = Spheres( + self.marker_trajectory[:, mi, :][:, np.newaxis, :], + color=color, + radius=radius, + name=marker_name, + **kwargs, + ) + else: + markers_seq = PointClouds( + self.marker_trajectory[:, mi, :][:, np.newaxis, :], + name=marker_name, + point_size=point_size, + color=color, + **kwargs, + ) + markers_seq.enabled = False + self._add_node(markers_seq) + + @classmethod + def from_c3d( + cls, + c3d_path, + start_frame=None, + end_frame=None, + fps_out=None, + colors=None, + lengths=None, + nb_markers_expected=None, + point_size=5.0, + color=(0.0, 0.0, 1.0, 1.0), + y_up=True, + **kwargs, + ): + """Load a sequence from an npz file. The filename becomes the name of the sequence""" + + markers_array, markers_labels, fps_in = load_markers(c3d_path, nb_markers_expected) + print(f"fps_in={fps_in} fps_out={fps_out} markers_array.shape={markers_array.shape}") + + if y_up: + markers_array[:, :, [0, 1, 2]] = markers_array[:, :, [0, 2, 1]] # Swap y and z + markers_array[:, :, 2] = -markers_array[:, :, 2] # Flip z + + # print(markers_array) + name = "Mocap " + os.path.splitext(os.path.basename(c3d_path))[0] + + # Crop frames and resample + sf = start_frame or 0 + ef = end_frame or markers_array.shape[0] + markers_array = markers_array[sf:ef] + + if fps_out is not None and fps_in != fps_out: + assert fps_in % fps_out == 0, "fps_out must be a interger divisor of fps_in" + mask = np.arange(0, markers_array.shape[0], fps_in // fps_out) + # markers_array = resample_positions(markers_array, fps_in, fps_out) # This uses splines and don't deal with NaN + markers_array = markers_array[mask] + + return cls( + markers_labels=markers_labels, + points=markers_array, + name=name, + colors=colors, + lengths=lengths, + point_size=point_size, + color=color, + **kwargs, + ) + + @classmethod + def from_synthetic( + cls, + synth_mocap_path, + start_frame=None, + end_frame=None, + fps_out=None, + colors=None, + lengths=None, + nb_markers_expected=None, + point_size=5.0, + color=(0.0, 0.0, 1.0, 1.0), + **kwargs, + ): + """Load a sequence from an npz file. The filename becomes the name of the sequence""" + + assert os.path.exists(synth_mocap_path), f"File {synth_mocap_path} does not exist" + # Load the marker trajectories + synthetic_markers = pkl.load(open(synth_mocap_path, "rb")) + + fps_in = int(synthetic_markers.fps) + assert 1 - int(synthetic_markers.fps) / fps_in < 1e-3, "fps must be an integer" + markers_labels = synthetic_markers.marker_names + name = "Synthetic markers" + + markers_array = synthetic_markers.marker_trajectory + + if fps_out is not None and abs(1 - fps_in / fps_out) > 1e-4: + assert ( + fps_in % fps_out == 0 + ), f"fps_out must be a interger divisor of fps_in, but got fps_in={fps_in} fps_out={fps_out}" + mask = np.arange(0, markers_array.shape[0], int(fps_in // fps_out)) + + # markers_array = resample_positions(markers_array, fps_in, fps_out) # This uses splines and don't deal with NaN + markers_array = markers_array[mask] + + return cls( + markers_labels=markers_labels, + points=markers_array, + name=name, + colors=colors, + lengths=lengths, + point_size=point_size, + color=color, + **kwargs, + ) + + @classmethod + def from_SSM_pkl( + cls, ssm_pkl_path, fps_out=None, colors=None, lengths=None, point_size=5.0, color=(0.0, 0.0, 1.0, 1.0), **kwargs + ): + """Load a sequence from an npz file. The filename becomes the name of the sequence""" + + # Load the marker trajectories + markers_data = pkl.load( + open(ssm_pkl_path, "rb"), encoding="latin1" + ) # dict_keys(['labels', 'required_parameters', 'markers']) + + fps_in = int(markers_data["required_parameters"]["frame_rate"]) + fps_in = ( + 60 # For the SSM dataset, the frame rate specified in the pkl file appears wrong, I assume it is 60 fps + ) + # assert abs(1-fps_in/markers_data['required_parameters']['frame_rate'])<1e-3, 'Frame rate is not an integer' + + markers_labels = [label.decode("utf-8") for label in markers_data["labels"]] + name = "SSM markers" + + markers_array = markers_data["markers"] + + # rotate the mocap data to align them with amass + markers_array[:, :, [0, 1, 2]] = markers_array[:, :, [2, 1, 0]] # Swap x and z + markers_array[:, :, 0] = -markers_array[:, :, 0] # Flip x + + if fps_out is not None and fps_in != fps_out: + assert fps_in % fps_out == 0, "fps_out must be a interger divisor of fps_in" + mask = np.arange(0, markers_array.shape[0], int(fps_in // fps_out)) + + # markers_array = resample_positions(markers_array, fps_in, fps_out) # This uses splines and don't deal with NaN + markers_array = markers_array[mask] + + print(f"fps_in={fps_in} fps_out={fps_out} markers_array.shape={markers_array.shape}") + + return cls( + markers_labels=markers_labels, + points=markers_array, + name=name, + colors=colors, + lengths=lengths, + point_size=point_size, + color=color, + **kwargs, + ) + + @classmethod + def from_file(cls, mocap_file, **kwargs): + if mocap_file.endswith(".c3d"): + return cls.from_c3d(mocap_file, **kwargs) + elif mocap_file.endswith(".npz"): + return cls.from_synthetic(mocap_file, **kwargs) + elif mocap_file.endswith(".pkl"): + return cls.from_SSM_pkl(mocap_file, **kwargs) + else: + raise ValueError(f"Unknown mocap file format: {mocap_file}, must be .c3d, .npz or .pkl") diff --git a/aitviewer/renderables/osim.py b/aitviewer/renderables/osim.py new file mode 100644 index 00000000..f49d74b0 --- /dev/null +++ b/aitviewer/renderables/osim.py @@ -0,0 +1,487 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import os +import pickle as pkl +import shutil + +import numpy as np +import tqdm +import trimesh + +from aitviewer.configuration import CONFIG as C +from aitviewer.renderables.markers import Markers +from aitviewer.renderables.meshes import Meshes +from aitviewer.renderables.rigid_bodies import RigidBodies +from aitviewer.scene.node import Node +from aitviewer.utils import to_numpy as c2c +from aitviewer.utils.colors import vertex_colors_from_weights + +try: + import nimblephysics as nimble +except ImportError: + raise ImportError("nimblephysics not found. Please install nimblephysics to use this module.") + + +def load_osim(osim_path, geometry_path=os.path.join(C.skel_models, "Geometry"), ignore_geometry=False): + """Load an osim file""" + + assert os.path.exists(osim_path), f"Could not find osim file {os.path.abspath(osim_path)}" + osim_path = os.path.abspath(osim_path) + + # Check that there is a Geometry folder at the same level as the osim file + file_geometry_path = os.path.join(os.path.dirname(osim_path), "Geometry") + + if not os.path.exists(file_geometry_path) or ignore_geometry: + if ignore_geometry and os.path.exists(file_geometry_path): + print(f"WARNING: Ignoring geometry folder at {file_geometry_path}") + else: + print(f"WARNING: No Geometry folder found at {file_geometry_path}, using {geometry_path} instead") + # Create a copy of the osim file at the same level as the geometry folder + tmp_osim_file = os.path.join(geometry_path, "..", "tmp.osim") + if os.path.exists(tmp_osim_file): + # remove the old file + os.remove(tmp_osim_file) + shutil.copyfile(osim_path, tmp_osim_file) + print(f"Copied {osim_path} to {tmp_osim_file}") + osim_path = os.path.abspath(tmp_osim_file) + + osim: nimble.biomechanics.OpenSimFile = nimble.biomechanics.OpenSimParser.parseOsim(osim_path) + assert osim is not None, "Could not load osim file: {}".format(osim_path) + return osim + + +class OSIMSequence(Node): + """ + Represents a temporal sequence of OSSO poses. Can be loaded from disk or initialized from memory. + """ + + def __init__( + self, + osim, + motion, + color_markers_per_part=False, + color_markers_per_index=False, # Overrides color_markers_per_part + color_skeleton_per_part=False, + osim_path=None, + fps=None, + fps_in=None, + is_rigged=False, + show_joint_angles=False, + viewer=True, + z_up=False, + **kwargs, + ): + """ + Initializer. + :param osim_path: A osim model + :param mot: A motion array + :osim_path: Path the osim model was loaded from (optional) + :param kwargs: Remaining arguments for rendering. + """ + self.osim_path = osim_path + self.osim = osim + self.motion = motion + + assert self.osim_path, "No osim path given" + + self.fps = fps + self.fps_in = fps_in + + self._show_joint_angles = show_joint_angles + self._is_rigged = is_rigged or show_joint_angles + self._z_up = z_up + + assert len(motion.shape) == 2 + super(OSIMSequence, self).__init__(n_frames=motion.shape[0], **kwargs) + + self._render_kwargs = kwargs + + # The node names of the skeleton model, the associated mesh and the template indices + body_nodes = [osim.skeleton.getBodyNode(i) for i in range(osim.skeleton.getNumBodyNodes())] + self.node_names = [n.getName() for n in body_nodes] + + self.meshes_dict = {} + self.indices_dict = {} + self.generate_meshes_dict() # Populate self.meshes_dict and self.indices_dict + self.create_template() + + # model markers + markers_labels = [ml for ml in self.osim.markersMap.keys()] + markers_labels.sort() + self.markers_labels = markers_labels + + # Nodes + self.vertices, self.faces, self.marker_trajectory, self.joints, self.joints_ori = self.fk() + + if viewer == False: + return + + if self._z_up and not C.z_up: + self.rotation = np.matmul(np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]), self.rotation) + + + if self._show_joint_angles: + global_oris = self.joints_ori + self.rbs = RigidBodies(self.joints, global_oris, length=0.01, name="Joint Angles") + self._add_node(self.rbs) + + # Add meshes + kwargs = self._render_kwargs.copy() + kwargs["name"] = "Mesh" + kwargs["color"] = kwargs.get("color", (160 / 255, 160 / 255, 160 / 255, 1.0)) + if color_skeleton_per_part: + kwargs["vertex_colors"] = self.per_part_bone_colors() + self.mesh_seq = Meshes(self.vertices, self.faces, **kwargs) + self._add_node(self.mesh_seq) + + # Add markers + kwargs = self._render_kwargs.copy() + kwargs["name"] = "Markers" + kwargs["color"] = kwargs.get("color", (0 / 255, 0 / 255, 255 / 255, 1.0)) + if color_markers_per_part: + markers_color = self.per_part_marker_colors() + kwargs["colors"] = markers_color + if color_markers_per_index: + marker_index_colors = vertex_colors_from_weights( + weights=range(len(self.marker_trajectory[0])), scale_to_range_1=True, alpha=1 + )[np.newaxis, :, :] + marker_index_colors = list(marker_index_colors) + markers_color = marker_index_colors + import ipdb + + ipdb.set_trace() + kwargs["colors"] = marker_index_colors + self.markers_seq = Markers( + points=self.marker_trajectory, markers_labels=self.markers_labels, point_size=10.0, **kwargs + ) + self._add_node(self.markers_seq) + + def color_by_vertex_id(self): + """ + Color the mesh by vertex index. + """ + self.mesh_seq.color_by_vertex_index() + + def per_part_bone_colors(self): + """Color the mesh with one color per node.""" + vertex_colors = np.ones((self.n_frames, self.template.vertices.shape[0], 4)) + color_palette = vertex_colors_from_weights(np.arange(len(self.node_names)), shuffle=True) + for i, node_name in enumerate(self.node_names): + id_start, id_end = self.indices_dict[node_name] + vertex_colors[:, id_start:id_end, 0:3] = color_palette[i, :] + return vertex_colors + + def per_part_marker_colors(self): + colors = vertex_colors_from_weights(np.arange(len(self.node_names)), alpha=1, shuffle=True) + + # Try to load a saved rigging file + rigging_file = None + if self.osim_path is not None: + # try to load a rigging file + rigging_file = self.osim_path.replace(".osim", f"_rigging.pkl") + + if not rigging_file is None and os.path.exists(rigging_file): + print(f"Loading rigging file from {rigging_file}") + rigging = pkl.load(open(rigging_file, "rb")) + marker_colors = colors[rigging["per_marker_rigging"]] + + else: + print(f"No rigging file {rigging_file} found. Fetching rigging for coloring.") + colors = vertex_colors_from_weights(np.arange(len(self.node_names)), alpha=1, shuffle=True) + + markers_rigging = 1 * -np.ones(self.marker_trajectory.shape[1]) + marker_colors = np.ones((self.marker_trajectory.shape[1], 4)) + + for mi, ml in (pbar := tqdm.tqdm(enumerate(self.markers_labels))): + pbar.set_description("Computing the per marker rigging ") + bone = self.osim.markersMap[ml][0].getName() + bone_index = self.node_names.index(bone) + markers_rigging[mi] = bone_index + color = colors[bone_index] + marker_colors[mi] = color + # print(marker_colors) + + return marker_colors + + def generate_meshes_dict(self): + """Output a dictionary giving for each bone, the attached mesh""" + + current_index = 0 + self.indices_dict = {} + self.meshes_dict = {} + + node_names = self.node_names + for node_name in node_names: + mesh_list = [] + body_node = self.osim.skeleton.getBodyNode(node_name) + # print(f' Loading meshes for node: {node_name}') + num_shape_nodes = body_node.getNumShapeNodes() + if num_shape_nodes == 0: + print(f"WARNING:\tNo shape nodes listed for {node_name}") + for shape_node_i in range(num_shape_nodes): + shape_node = body_node.getShapeNode(shape_node_i) + submesh_path = shape_node.getShape().getMeshPath() + # Get the scaling for this meshes + scale = shape_node.getShape().getScale() + offset = shape_node.getRelativeTranslation() + # Load the mesh + try: + submesh = trimesh.load_mesh(submesh_path, process=False) + # print(f'Loaded mesh {submesh_path}') + except Exception as e: + print(e) + print(f"WARNING:\tCould not load mesh {submesh_path}") + submesh = None + continue + + if submesh is not None: + trimesh.repair.fix_normals(submesh) + trimesh.repair.fix_inversion(submesh) + trimesh.repair.fix_winding(submesh) + + # import pyvista + # submesh_poly = pyvista.read(submesh_path) + # faces_as_array = submesh_poly.faces.reshape((submesh_poly.n_faces, 4))[:, 1:] + # submesh = trimesh.Trimesh(submesh_poly.points, faces_as_array) + + # Scale the bone to match .osim subject scaling + submesh.vertices[:] = submesh.vertices * scale + submesh.vertices[:] += offset + # print(f'submesh_path: {submesh_path}, Nb vertices: {submesh.vertices.shape[0]}') + mesh_list.append(submesh) + + # Concatenate meshes + if mesh_list: + node_mesh = trimesh.util.concatenate(mesh_list) + self.indices_dict[node_name] = (current_index, current_index + node_mesh.vertices.shape[0]) + current_index += node_mesh.vertices.shape[0] + else: + node_mesh = None + print("\t WARNING: No submesh for node:", node_name) + self.indices_dict[node_name] = (current_index, current_index) + + # Add to the dictionary + self.meshes_dict[node_name] = node_mesh + print(self.meshes_dict) + + def create_template(self): + part_meshes = [] + for node_name in self.node_names: + mesh = self.meshes_dict[node_name] + # assert mesh, "No mesh for node: {}".format(node_name) + if mesh is None: + print("WARNING: No mesh for node: {}".format(node_name)) + if mesh: + part_meshes.append(mesh) + # part_meshes = [m for m in part_meshes if m] + template = trimesh.util.concatenate(part_meshes) + # import ipdb; ipdb.set_trace() + + template.remove_degenerate_faces() + self.template = template + + # save mesh + # # import ipdb; ipdb.set_trace() + # self.template.export('template.obj') + # print(f'Saved template to template.obj') + + # from psbody.mesh import Mesh + # m = Mesh(filename='template.obj') + # m.set_vertex_colors_from_weights(np.arange(m.v.shape[0])) + # m.show() + + @classmethod + def a_pose(cls, osim_path=None, **kwargs): + """Creates a OSIM sequence whose single frame is a OSIM mesh in rest pose.""" + # Load osim file + if osim_path is None: + # osim: nimble.biomechanics.OpenSimFile = nimble.models.RajagopalHumanBodyModel() + # osim_path = "RajagopalHumanBodyModel.osim" # This is not a real path, but it is needed to instantiate the sequence object + osim_path = os.path.join(C.skel_models, "bsm.osim") + osim = load_osim(osim_path) + else: + osim = load_osim(osim_path) + + assert osim is not None, "Could not load osim file: {}".format(osim_path) + motion = osim.skeleton.getPositions()[np.newaxis, :] + + return cls(osim, motion, osim_path=osim_path, **kwargs) + + @classmethod + def zero_pose(cls, osim_path=None, **kwargs): + """Creates a OSIM sequence whose single frame is a OSIM mesh in rest pose.""" + # Load osim file + if osim_path is None: + osim: nimble.biomechanics.OpenSimFile = nimble.models.RajagopalHumanBodyModel() + osim_path = "RajagopalHumanBodyModel.osim" # This is not a real path, but it is needed to instantiate the sequence object + else: + osim = nimble.biomechanics.OpenSimParser.parseOsim(osim_path) + + assert osim is not None, "Could not load osim file: {}".format(osim_path) + + # motion = np.zeros((1, len(osim.skeleton.getBodyNodes()))) + motion = osim.skeleton.getPositions()[np.newaxis, :] + motion = np.zeros_like(motion) + # import ipdb; ipdb.set_trace() + + return cls(osim, motion, osim_path=osim_path, **kwargs) + + @classmethod + def from_ab_folder(cls, ab_folder, trial, start_frame=None, end_frame=None, fps_out=None, **kwargs): + """ + Load an osim sequence from a folder returned by AddBiomechanics + ab_folder: the folder returned by AddBiomechanics, ex: '/home/kellerm/Data/AddBiomechanics/CMU/01/smpl_head_manual' + trial: Trial name + start_frame: the first frame to load + end_frame: the last frame to load + fps_out: the output fps + """ + + if ab_folder[-1] != "/": + ab_folder += "/" + + mot_file = ab_folder + f"IK/{trial}_ik.mot" + osim_path = ab_folder + "Models/optimized_scale_and_markers.osim" + + return OSIMSequence.from_files( + osim_path=osim_path, + mot_file=mot_file, + start_frame=start_frame, + end_frame=end_frame, + fps_out=fps_out, + **kwargs, + ) + + @classmethod + def from_files( + cls, + osim_path, + mot_file, + start_frame=None, + end_frame=None, + fps_out: int = None, + ignore_fps=False, + ignore_geometry=False, + **kwargs, + ): + """Creates a OSIM sequence from addbiomechanics return data + osim_path: .osim file path + mot_file : .mot file path + start_frame: first frame to use in the sequence + end_frame: last frame to use in the sequence + fps_out: frames per second of the output sequence + ignore_geometry : use the aitconfig.osim_geometry folder instead of the one next to the osim file + """ + + # Nimblephysics does not like relative paths + osim_path = os.path.abspath(osim_path) + mot_file = os.path.abspath(mot_file) + + # Load osim file + osim = load_osim(osim_path, ignore_geometry=ignore_geometry) + + # Load the .mot file + mot: nimble.biomechanics.OpenSimMot = nimble.biomechanics.OpenSimParser.loadMot(osim.skeleton, mot_file) + + motion = np.array(mot.poses.T) + + # Crop and sample + sf = start_frame or 0 + ef = end_frame or motion.shape[0] + motion = motion[sf:ef] + + # estimate fps_in + ts = np.array(mot.timestamps) + fps_estimated = 1 / np.mean(ts[1:] - ts[:-1]) + fps_in = int(round(fps_estimated)) + print(f"Estimated fps for the .mot sequence: {fps_estimated}, rounded to {fps_in}") + + if not ignore_fps: + assert ( + abs(1 - fps_estimated / fps_in) < 1e-5 + ), f"FPS estimation might be bad, {fps_estimated} rounded to {fps_in}, check." + + if fps_out is not None: + assert fps_in % fps_out == 0, "fps_out must be a interger divisor of fps_in" + mask = np.arange(0, motion.shape[0], fps_in // fps_out) + print(f"Resampling from {fps_in} to {fps_out} fps. Keeping every {fps_in//fps_out}th frame") + # motion = resample_positions(motion, fps_in, fps_out) #TODO: restore this + motion = motion[mask] + + del mot + else: + fps_out = fps_in + + return cls(osim, motion, osim_path=osim_path, fps=fps_out, fps_in=fps_in, **kwargs) + + def fk(self): + """Get vertices from the poses.""" + # Forward kinematics https://github.com/nimblephysics/nimblephysics/search?q=setPositions + + verts = np.zeros((self.n_frames, self.template.vertices.shape[0], self.template.vertices.shape[1])) + markers = np.zeros((self.n_frames, len(self.markers_labels), 3)) + + joints = np.zeros([self.n_frames, len(self.meshes_dict), 3]) + joints_ori = np.zeros([self.n_frames, len(self.meshes_dict), 3, 3]) + + prev_verts = verts[0] + prev_pose = self.motion[0, :] + + for frame_id in (pbar := tqdm.tqdm(range(self.n_frames))): + pbar.set_description("Generating osim skeleton meshes ") + + pose = self.motion[frame_id, :] + # If the pose did not change, use the previous frame verts + if np.all(pose == prev_pose) and frame_id != 0: + verts[frame_id] = prev_verts + continue + + # Pose osim + self.osim.skeleton.setPositions(self.motion[frame_id, :]) + + # Since python 3.6, dicts have a fixed order so the order of this list should be marching labels + markers[frame_id, :, :] = np.vstack( + self.osim.skeleton.getMarkerMapWorldPositions(self.osim.markersMap).values() + ) + # Sanity check for previous comment + assert ( + list(self.osim.skeleton.getMarkerMapWorldPositions(self.osim.markersMap).keys()) == self.markers_labels + ), "Marker labels are not in the same order" + + for ni, node_name in enumerate(self.node_names): + # if ('thorax' in node_name) or ('lumbar' in node_name): + # # We do not display the spine as the riggidly rigged mesh can't represent the constant curvature of the spine + # continue + mesh = self.meshes_dict[node_name] + if mesh is not None: + part_verts = mesh.vertices + + # pose part + transfo = self.osim.skeleton.getBodyNode(node_name).getWorldTransform() + + # Add a row of homogenous coordinates + part_verts = np.concatenate([part_verts, np.ones((mesh.vertices.shape[0], 1))], axis=1) + part_verts = np.matmul(part_verts, transfo.matrix().T)[:, 0:3] + + # Update the part in the full mesh + id_start, id_end = self.indices_dict[node_name] + verts[frame_id, id_start:id_end, :] = part_verts + + # Update joint + joints[frame_id, ni, :] = transfo.translation() + joints_ori[frame_id, ni, :, :] = transfo.rotation() + + prev_verts = verts[frame_id] + prev_pose = pose + + faces = self.template.faces + + return c2c(verts), c2c(faces), markers, joints, joints_ori + + def redraw(self): + self.vertices, self.faces, self.marker_trajectory, self.joints, self.joints_ori = self.fk() + if self._is_rigged: + self.skeleton_seq.joint_positions = self.joints + self.mesh_seq.vertices = self.vertices + self.marker_seq = self.marker_trajectory + super().redraw() diff --git a/aitviewer/renderables/skel.py b/aitviewer/renderables/skel.py new file mode 100644 index 00000000..09e75031 --- /dev/null +++ b/aitviewer/renderables/skel.py @@ -0,0 +1,731 @@ +""" +Copyright©2023 Max-Planck-Gesellschaft zur Förderung +der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute +for Intelligent Systems. All rights reserved. + +Author: Marilyn Keller +See https://skel.is.tue.mpg.de/license.html for licensing and contact information. +""" + +import os +import pickle as pkl + +import numpy as np +import torch +import tqdm +import trimesh +from scipy.spatial.transform import Rotation + +from aitviewer.configuration import CONFIG as C +from aitviewer.models.smpl import SMPLLayer +from aitviewer.renderables.meshes import Meshes +from aitviewer.renderables.rigid_bodies import RigidBodies +from aitviewer.renderables.skeletons import Skeletons +from aitviewer.scene.node import Node +from aitviewer.utils import interpolate_positions, local_to_global, resample_positions +from aitviewer.utils import to_numpy as c2c +from aitviewer.utils import to_torch +from aitviewer.utils.decorators import hooked +from aitviewer.utils.so3 import aa2euler_numpy +from aitviewer.utils.so3 import aa2rot_torch as aa2rot +from aitviewer.utils.so3 import ( + euler2aa_numpy, + interpolate_rotations, + resample_rotations, +) +from aitviewer.utils.so3 import rot2aa_torch as rot2aa + +try: + from skel.kin_skel import skel_joints_name + from skel.skel_model import SKEL +except ImportError as e: + raise ImportError(f"Could not import SKEL. Please install it from https://github.com/MarilynKeller/skel.git") + +from aitviewer.utils.colors import skining_weights_to_color, vertex_colors_from_weights + + +class SKELSequence(Node): + """ + Represents a temporal sequence of SMPL poses. Can be loaded from disk or initialized from memory. + """ + + def __init__( + self, + poses_body, + skel_layer, + poses_type="skel", + betas=None, + trans=None, + device=C.device, + dtype=C.f_precision, + include_root=True, + normalize_root=False, + is_rigged=False, + show_joint_angles=False, + z_up=False, + fps=None, + fps_in=None, + show_joint_arrows=True, + visual=True, + post_fk_func=None, + skin_color=(200 / 255, 160 / 255, 160 / 255, 125 / 255), + skel_color=(160 / 255, 160 / 255, 160 / 255, 255 / 255), + skin_coloring=None, # "pose_offsets", "skinning_weights" + skel_coloring=None, # "skinning_weights", "bone_label" + **kwargs, + ): + """ + Initializer. + :param poses_body: An array (numpy ar pytorch) of shape (F, 46) containing the pose parameters of the + body, i.e. without hands or face parameters. + :skel_layer: The SKEL layer that maps parameters to joint positions and/or body and meshes surfaces. + :param betas: An array (numpy or pytorch) of shape (N_BETAS, ) containing the shape parameters. + :param trans: An array (numpy or pytorch) of shape (F, 3) containing a global translation that is applied to + all joints and vertices. + :param device: The pytorch device for computations. + :param dtype: The pytorch data type. + :param include_root: Whether or not to include root information. If False, no root translation and no root + rotation is applied. + :param normalize_root: Whether or not to normalize the root. If True, the global root translation in the first + frame is zero and the global root orientation is the identity. + :param is_rigged: Whether or not to display the joints as a skeleton. + :param show_joint_angles: Whether or not the coordinate frames at the joints should be visualized. + :param z_up: Whether or not the input data assumes Z is up. If so, the data will be rotated such that Y is up. + :param post_fk_func: User specified postprocessing function that is called after evaluating the SMPL model, + the function signature must be: def post_fk_func(self, vertices, joints, current_frame_only), + and it must return new values for vertices and joints with the same shapes. + Shapes are: + if current_frame_only is False: vertices (F, V, 3) and joints (F, N_JOINTS, 3) + if current_frame_only is True: vertices (1, V, 3) and joints (1, N_JOINTS, 3) + :param skin_coloring: Coloring the skin mesh of SKEL per vertex. Must be in ['skinning_weights', 'pose_offsets']. + :param skel_coloring: Coloring the bones mesh of SKEL per vertex. Must be in ['skinning_weights', 'bone_label']. + :param kwargs: Remaining arguments for rendering. + """ + assert len(poses_body.shape) == 2 + + super(SKELSequence, self).__init__(n_frames=poses_body.shape[0], **kwargs) + self.skel_layer = skel_layer + self.post_fk_func = post_fk_func + + self.device = device + self.fps = fps # fps of this loaded sequence + self.fps_in = fps_in # original fps of the sequence + + self.poses_body = to_torch(poses_body, dtype=dtype, device=device) + self.poses_type = poses_type + + betas = betas if betas is not None else torch.zeros([1, self.skel_layer.num_betas]) + trans = trans if trans is not None else torch.zeros([len(poses_body), 3]) + + self.betas = to_torch(betas, dtype=dtype, device=device) + self.trans = to_torch(trans, dtype=dtype, device=device) + + if len(self.betas.shape) == 1: + self.betas = self.betas.unsqueeze(0) + + self._include_root = include_root + self._normalize_root = normalize_root + self._show_joint_angles = show_joint_angles + self._is_rigged = is_rigged or show_joint_angles + self._render_kwargs = kwargs + self._z_up = z_up + + if not self._include_root: + self.trans = torch.zeros_like(self.trans) + + # Edit mode + self.gui_modes.update({"edit": {"title": " Edit", "fn": self.gui_mode_edit, "icon": "\u0081"}}) + + self._edit_joint = None + self._edit_pose = None + self._edit_pose_dirty = False + + # Nodes + skel_output = self.fk() + + self.skin_vertices = skel_output.skin_verts + self.skel_vertices = skel_output.skel_verts + self.skin_faces = skel_output.skin_f + self.skel_faces = skel_output.skel_f + self.joints = skel_output.joints + self.joints_ori = skel_output.joints_ori + self.skeleton = skel_output.skeleton + + self.skel_output = skel_output + + if self._z_up: + self.rotation = np.matmul(np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]), self.rotation) + + if visual == False: + return + + if self._is_rigged: + # Must first add skeleton, otherwise transparency does not work correctly. + # Overriding given color with a custom color for the skeleton. + kwargs = self._render_kwargs.copy() + color = (1.0, 177 / 255, 1 / 255, 1.0) + + self.skeleton_seq = Skeletons(self.joints, self.skeleton, gui_affine=False, color=color, name="Kin tree") + # self.skeleton_seq.position = self.position + # self.skeleton_seq.rotation = self.rotation + self._add_node(self.skeleton_seq) + + global_oris = self.joints_ori + + if show_joint_arrows is False: + arrow_length = 0 + else: + arrow_length = 0.1 + + self.rbs = RigidBodies( + self.joints, global_oris, length=arrow_length, name="Joint Angles", color=(1.0, 177 / 255, 1 / 255, 1.0) + ) + self._add_node(self.rbs, enabled=self._show_joint_angles) + + # Instantiate the Skin submesh with proper colouring + kwargs = self._render_kwargs.copy() + mesh_name = "Skin" + skin_colors = None + if skin_coloring == "skinning_weights": + skin_colors = skining_weights_to_color( + skel_layer.skin_weights.to_dense().cpu().numpy(), alpha=skin_color[-1] + ) + elif skin_coloring == "pose_offsets": + values = self.skel_output.pose_offsets.cpu().numpy() + skin_colors = values / np.max(np.abs(values)) + # append an alpha channel + skin_colors = np.concatenate( + [skin_colors, np.ones([skin_colors.shape[0], skin_colors.shape[1], 1])], axis=-1 + ) + self.skin_mesh_seq = Meshes( + self.skin_vertices, self.skin_faces, gui_affine=False, is_selectable=False, vertex_colors=skin_colors + ) + + if skin_colors is None: + self.skin_mesh_seq = Meshes( + self.skin_vertices, + self.skin_faces, + gui_affine=False, + is_selectable=False, + color=skin_color, + name=mesh_name, + ) + else: + self.skin_mesh_seq = Meshes( + self.skin_vertices, + self.skin_faces, + gui_affine=False, + is_selectable=False, + vertex_colors=skin_colors, + name=mesh_name, + ) + + self._add_node(self.skin_mesh_seq) + + # Instantiate the Bones submesh with proper colouring + mesh_name = "Bones" + skel_colors = None + if skel_coloring == "skinning_weights": + skel_colors = skining_weights_to_color(skel_layer.skel_weights.cpu().numpy(), alpha=255) + elif skel_coloring == "bone_label": + skel_colors = skining_weights_to_color(skel_layer.skel_weights_rigid.cpu().numpy(), alpha=255) + + if skel_colors is None: + self.bones_mesh_seq = Meshes( + self.skel_vertices, + self.skel_faces, + gui_affine=False, + is_selectable=False, + color=skel_color, + name=mesh_name, + ) + else: + draw_bones_outline = False + if skin_coloring == "skinning_weights" and ( + skel_coloring == "skinning_weights" or skel_coloring == "bone_label" + ): + draw_bones_outline = True # For better visibility of the bones + self.bones_mesh_seq = Meshes( + self.skel_vertices, + self.skel_faces, + gui_affine=False, + is_selectable=False, + vertex_colors=skel_colors, + name=mesh_name, + draw_outline=draw_bones_outline, + ) + + self._add_node(self.bones_mesh_seq) + + # Save view mode state to restore when exiting edit mode. + self._skin_view_mode_color = self.skin_mesh_seq.color + self._skel_view_mode_color = self.bones_mesh_seq.color + self._view_mode_joint_angles = self._show_joint_angles + + def get_rotated_global_joint(self): + rot_smpl_joints = np.matmul(self.joints, self.rotation.T) + + rot_joints_ori = np.zeros_like(self.joints_ori) + for joint_idx in range(self.joints_ori.shape[1]): + rot_joints_ori[:, joint_idx, :, :] = np.matmul(self.rotation, self.joints_ori[:, joint_idx, :, :]) + + return rot_smpl_joints, rot_joints_ori + + @property + def rotated_vertices(self): + return np.matmul(self.vertices, self.rotation.T) + self.position + + @property + def rotated_skel_vertices(self): + return np.matmul(self.skel_vertices, self.rotation.T) + self.position + + @classmethod + def from_file( + cls, + skel_seq_file, + fps_in, + start_frame=None, + end_frame=None, + log=True, + fps_out=None, + z_up=False, + device=C.device, + poses_type="skel", + **kwargs, + ): + """Load a SKEL sequence from a pkl.""" + + if skel_seq_file.endswith(".pkl"): + skel_data = pkl.load(open(skel_seq_file, "rb")) + elif skel_seq_file.endswith(".npz"): + # Compatibility with PS fitting pipeline + skel_data = np.load(skel_seq_file) + skel_data = {key: skel_data[key] for key in skel_data.files} + if "poses" not in skel_data and "pose" in skel_data and "global_orient" in skel_data: + skel_data["poses"] = np.concatenate( + [ + skel_data["global_orient"], + skel_data["pose"], + ], + axis=1, + ) + if "trans" not in skel_data and "transl" in skel_data: + skel_data["trans"] = skel_data["transl"] + del skel_data["transl"] + if "gender" not in skel_data: + print("Warning: no gender found in the npz file, assuming female.") + skel_data["gender"] = "female" + if skel_data["betas"].shape[0] == 1: + skel_data["betas"] = skel_data["betas"].repeat(skel_data["poses"].shape[0], axis=0) + + import ipdb + + ipdb.set_trace() + + else: + raise ValueError(f"skel_seq_file must be a pkl or npz file, got {skel_seq_file}") + + for key in ["poses", "trans", "betas", "gender"]: + assert ( + key in skel_data + ), f"The loaded skel sequence dictionary must contain {key}. Loaded dictionary has keys: {skel_data.keys()}" + + gender = skel_data["gender"] + skel_layer = SKEL(model_path=C.skel_models, gender=gender) + + assert gender == skel_layer.gender, f"skel layer has gender {skel_layer.gender} while data has gender {gender}" + + sf = start_frame or 0 + ef = end_frame or skel_data["poses"].shape[0] + poses = skel_data["poses"][sf:ef] + trans = skel_data["trans"][sf:ef] + betas = skel_data["betas"][sf:ef] + + if fps_out is not None: + if fps_in != fps_out: + betas = resample_positions(betas, fps_in, fps_out) + poses = resample_positions(poses, fps_in, fps_out) # Linear interpolation + print("WARNING: poses resampled with linear interpolation, this is wrong but ok for int fps ratio") + trans = resample_positions(trans, fps_in, fps_out) + else: + fps_out = fps_in + + i_beta_end = skel_layer.num_betas + return cls( + poses_body=poses, + skel_layer=skel_layer, + betas=betas[:, :i_beta_end], + trans=trans, + z_up=z_up, + device=device, + fps=fps_out, + fps_in=fps_in, + poses_type=poses_type, + **kwargs, + ) + + @classmethod + def t_pose( + cls, skel_layer, betas=None, frames=1, is_rigged=True, show_joint_angles=False, device=C.device, **kwargs + ): + """Creates a SKEL sequence whose single frame is a SKEL mesh in T-Pose.""" + + if betas is not None: + assert betas.shape[0] == 1 + + poses = np.zeros([frames, skel_layer.num_q_params]) + return cls( + poses, + skel_layer=skel_layer, + betas=betas, + is_rigged=is_rigged, + show_joint_angles=show_joint_angles, + device=device, + **kwargs, + ) + + @property + def color(self): + return self.mesh_seq.color + + @color.setter + def color(self, color): + self.mesh_seq.color = color + + @property + def bounds(self): + return self.skin_mesh_seq.bounds + + @property + def current_bounds(self): + return self.skin_mesh_seq.current_bounds + + @property + def vertex_normals(self): + return self.skin_mesh_seq.vertex_normals + + @property + def poses(self): + return self.poses_body + + @property + def _edit_mode(self): + return self.selected_mode == "edit" + + def fk(self, current_frame_only=False): + """Get joints and/or vertices from the poses.""" + if current_frame_only: + # Use current frame data. + if self._edit_mode: + poses_body = self._edit_pose[None, :] + else: + poses_body = self.poses_body[self.current_frame_id][None, :] + + trans = self.trans[self.current_frame_id][None, :] + + if self.betas.shape[0] == self.n_frames: + betas = self.betas[self.current_frame_id][None, :] + else: + betas = self.betas + else: + # Use the whole sequence. + if self._edit_mode: + poses_body = self.poses_body.clone() + poses_body[self.current_frame_id] = self._edit_pose + else: + poses_body = self.poses_body + + trans = self.trans + betas = self.betas + + skel_output = self.skel_layer(poses=poses_body, betas=betas, trans=trans, poses_type=self.poses_type) + + # skin_verts = skel_output.skin_verts + # skel_verts = skel_output.skel_verts + # joints = skel_output.joints + # joints_ori = skel_output.joints_ori + + if current_frame_only: + # return c2c(skin_verts)[0], c2c(skin_f)[0], c2c(skel_verts)[0], c2c(skel_f)[0], c2c(joints)[0], c2c(joints_ori)[0], c2c(skeleton) + for att in ["skin_verts", "skel_verts", "joints", "joints_ori"]: + att_value = getattr(skel_output, att) + setattr(skel_output, att, c2c(att_value)[0]) + else: + # return c2c(skin_verts), c2c(skin_f), c2c(skel_verts), c2c(skel_f), c2c(joints), c2c(joints_ori), c2c(skeleton) + for att in ["skin_verts", "skel_verts", "joints", "joints_ori"]: + try: + att_value = getattr(skel_output, att) + setattr(skel_output, att, c2c(att_value)) + except: + import ipdb + + ipdb.set_trace() + + skel_output.skin_f = c2c(self.skel_layer.skin_f) + skel_output.skel_f = c2c(self.skel_layer.skel_f) + + skeleton = self.skel_layer.kintree_table.T + skeleton[0, 0] = -1 + skel_output.skeleton = c2c(skeleton) + + return skel_output + + def interpolate(self, frame_ids): + """ + Replace the frames at the given frame IDs via an interpolation of its neighbors. Only the body pose as well + as the root pose and translation are interpolated. + :param frame_ids: A list of frame ids to be interpolated. + """ + ids = np.unique(frame_ids) + all_ids = np.arange(self.n_frames) + mask_avail = np.ones(self.n_frames, dtype=np.bool) + mask_avail[ids] = False + + # Interpolate poses. + all_poses = torch.cat([self.poses_root, self.poses_body], dim=-1) + ps = np.reshape(all_poses.cpu().numpy(), (self.n_frames, -1, 3)) + ps_interp = interpolate_rotations(ps[mask_avail], all_ids[mask_avail], ids) + all_poses[ids] = torch.from_numpy(ps_interp.reshape(len(ids), -1)).to( + dtype=self.betas.dtype, device=self.betas.device + ) + self.poses_root = all_poses[:, :3] + self.poses_body = all_poses[:, 3:] + + # Interpolate global translation. + ts = self.trans.cpu().numpy() + ts_interp = interpolate_positions(ts[mask_avail], all_ids[mask_avail], ids) + self.trans[ids] = torch.from_numpy(ts_interp).to(dtype=self.betas.dtype, device=self.betas.device) + + self.redraw() + + @hooked + def on_before_frame_update(self): + if self._edit_mode and self._edit_pose_dirty: + self._edit_pose = self.poses[self.current_frame_id].clone() + self.redraw(current_frame_only=True) + self._edit_pose_dirty = False + + @hooked + def on_frame_update(self): + if self.edit_mode: + self._edit_pose = self.poses[self.current_frame_id].clone() + self._edit_pose_dirty = False + + def redraw(self, **kwargs): + current_frame_only = kwargs.get("current_frame_only", False) + + # Use the edited pose if in edit mode. + # skin_vertices, skin_faces, skel_vertices, skel_faces, joints, joints_ori, skeleton = self.fk(current_frame_only) + skel_output = self.fk(current_frame_only) + + if current_frame_only: + self.skin_vertices[self.current_frame_id] = skel_output.skin_verts + self.skel_vertices[self.current_frame_id] = skel_output.skel_verts + self.joints[self.current_frame_id] = skel_output.joints + + if self._is_rigged: + self.skeleton_seq.current_joint_positions = skel_output.joints + + # Use current frame data. + if self._edit_mode: + pose = self._edit_pose + else: + pose = self.poses_body[self.current_frame_id] + + # Update rigid bodies. + global_oris = skel_output.joints_ori + self.rbs.current_rb_ori = c2c(global_oris) + self.rbs.current_rb_pos = self.joints[self.current_frame_id] + + # Update mesh. + self.skin_mesh_seq.current_vertices = skel_output.skin_verts + self.bones_mesh_seq.current_vertices = skel_output.skel_verts + else: + self.skin_vertices = skel_output.skin_verts + self.skel_vertices = skel_output.skel_verts + self.joints = skel_output.joints + + # Update skeleton. + if self._is_rigged: + self.skeleton_seq.joint_positions = self.joints + + # Extract poses including the edited pose. + if self._edit_mode: + poses_body = self.poses_body.clone() + poses_body[self.current_frame_id] = self._edit_pose + else: + poses_body = self.poses_body + poses_root = self.poses_root + + # Update rigid bodies. + global_oris = skel_output.joints_ori + self.rbs.rb_ori = c2c(global_oris) + self.rbs.rb_pos = self.joints + + # Update mesh + self.skin_mesh_seq.vertices = skel_output.skin_verts + self.bones_mesh_seq.vertices = skel_output.skel_verts + + self.skel_output = skel_output + + super().redraw(**kwargs) + + @property + def edit_mode(self): + return self._edit_mode + + @property + def selected_mode(self): + return self._selected_mode + + @selected_mode.setter + def selected_mode(self, selected_mode): + if self._selected_mode == selected_mode: + return + self._selected_mode = selected_mode + + if self.selected_mode == "edit": + self.rbs.enabled = True + self.rbs.is_selectable = False + self._edit_pose = self.poses[self.current_frame_id].clone() + + # Disable picking for the mesh + self.skin_mesh_seq.backface_fragmap = True + self.bones_mesh_seq.backface_fragmap = True + self.rbs.color = (1, 0, 0.5, 1.0) + + self._skin_view_mode_color = self.skin_mesh_seq.color + self.skin_mesh_seq.color = (*self._skin_view_mode_color[:3], min(self._skin_view_mode_color[3], 0.5)) + + self._skel_view_mode_color = self.bones_mesh_seq.color + self.bones_mesh_seq.color = (*self._skel_view_mode_color[:3], min(self._skel_view_mode_color[3], 0.5)) + + self.redraw(current_frame_only=True) + + def _gui_joint(self, imgui, j, tree=None): + name = "unknown" + if j < len(skel_joints_name): + name = skel_joints_name[j] + + if tree: + e = imgui.tree_node(f"{j} - {name}") + else: + e = True + imgui.text(f"{j} - {name}") + + if e: + start_param = [0, 3, 6, 7, 8, 9, 10, 13, 14, 15, 16, 17, 20, 23, 26, 29, 32, 33, 34, 36, 39, 42, 43, 44, 46] + aa = self._edit_pose[start_param[j] : start_param[j + 1]] + if len(aa) == 1: + angle = np.degrees(aa.cpu().numpy()) + u, angle = imgui.drag_float(f"##joint{j}", angle, 0.1, format="%.2f") + if u: + aa = np.array(np.radians(angle)) + self._edit_pose[start_param[j] : start_param[j + 1]] = torch.from_numpy(aa) + self._edit_pose_dirty = True + self.redraw(current_frame_only=True) + elif len(aa) == 2: + angles = np.degrees(aa.cpu().numpy()) + u, angles = imgui.drag_float2(f"##joint{j}", *angles, 0.1, format="%.2f") + if u: + aa = np.radians(np.array(angles)) + self._edit_pose[start_param[j] : start_param[j + 1]] = torch.from_numpy(aa) + self._edit_pose_dirty = True + self.redraw(current_frame_only=True) + elif len(aa) == 3: + euler = aa2euler_numpy(aa.cpu().numpy(), degrees=True) + u, euler = imgui.drag_float3(f"##joint{j}", *euler, 0.1, format="%.2f") + if u: + aa = euler2aa_numpy(np.array(euler), degrees=True) + self._edit_pose[start_param[j] : start_param[j + 1]] = torch.from_numpy(aa) + self._edit_pose_dirty = True + self.redraw(current_frame_only=True) + if tree: + for c in tree.get(j, []): + self._gui_joint(imgui, c, tree) + imgui.tree_pop() + if tree: + for c in tree.get(j, []): + self._gui_joint(imgui, c, tree) + imgui.tree_pop() + + def gui_mode_edit(self, imgui): + kin_skel = self.skeleton + + tree = {} + for i in range(kin_skel.shape[1]): + if kin_skel[0, i] != -1: + tree.setdefault(kin_skel[0, i], []).append(kin_skel[1, i]) + + if not tree: + return + + if self._edit_joint is None: + self._gui_joint(imgui, 0, tree) + else: + self._gui_joint(imgui, self._edit_joint) + + if imgui.button("Apply"): + self.poses_root[self.current_frame_id] = self._edit_pose[:3] + self.poses_body[self.current_frame_id] = self._edit_pose[3:] + self._edit_pose_dirty = False + self.redraw(current_frame_only=True) + imgui.same_line() + if imgui.button("Apply to all"): + edit_rots = Rotation.from_rotvec(np.reshape(self._edit_pose.cpu().numpy(), (-1, 3))) + base_rots = Rotation.from_rotvec(np.reshape(self.poses[self.current_frame_id].cpu().numpy(), (-1, 3))) + relative = edit_rots * base_rots.inv() + for i in range(self.n_frames): + root = Rotation.from_rotvec(np.reshape(self.poses_root[i].cpu().numpy(), (-1, 3))) + self.poses_root[i] = torch.from_numpy((relative[0] * root).as_rotvec().flatten()) + + body = Rotation.from_rotvec(np.reshape(self.poses_body[i].cpu().numpy(), (-1, 3))) + self.poses_body[i] = torch.from_numpy((relative[1:] * body).as_rotvec().flatten()) + self._edit_pose_dirty = False + self.redraw() + imgui.same_line() + if imgui.button("Reset"): + self._edit_pose = self.poses[self.current_frame_id] + self._edit_pose_dirty = False + self.redraw(current_frame_only=True) + + def gui_io(self, imgui): + if imgui.button("Export sequence to NPZ"): + dir = os.path.join(C.export_dir, "SMPL") + os.makedirs(dir, exist_ok=True) + path = os.path.join(dir, self.name + ".npz") + self.export_to_npz(path) + print(f'Exported SMPL sequence to "{path}"') + + def gui_context_menu(self, imgui, x: int, y: int): + if self.edit_mode and self._edit_joint is not None: + self._gui_joint(imgui, self._edit_joint) + else: + if imgui.radio_button("View mode", not self.edit_mode): + self.selected_mode = "view" + imgui.close_current_popup() + if imgui.radio_button("Edit mode", self.edit_mode): + self.selected_mode = "edit" + imgui.close_current_popup() + + imgui.spacing() + imgui.separator() + imgui.spacing() + + super().gui_context_menu(imgui, x, y) + + def on_selection(self, node, instance_id, tri_id): + if self.edit_mode: + # Index of the joint that is currently being edited. + if node != self.skin_mesh_seq and node != self.bones_mesh_seq: + self._edit_joint = instance_id + self.rbs.color_one(self._edit_joint, (0.3, 0.4, 1, 1)) + else: + self._edit_joint = None + # Reset color of all spheres to the default color + self.rbs.color = self.rbs.color + + def render_outline(self, *args, **kwargs): + # Only render outline of the mesh, skipping skeleton and rigid bodies. + self.skin_mesh_seq.render_outline(*args, **kwargs) + self.bones_mesh_seq.render_outline(*args, **kwargs) diff --git a/aitviewer/renderables/smpl.py b/aitviewer/renderables/smpl.py index ad47430c..91ad4866 100644 --- a/aitviewer/renderables/smpl.py +++ b/aitviewer/renderables/smpl.py @@ -18,7 +18,9 @@ from aitviewer.utils import to_numpy as c2c from aitviewer.utils import to_torch from aitviewer.utils.decorators import hooked -from aitviewer.utils.so3 import aa2euler_numpy +from aitviewer.utils.so3 import ( + aa2euler_numpy, +) from aitviewer.utils.so3 import aa2rot_torch as aa2rot from aitviewer.utils.so3 import ( euler2aa_numpy, @@ -157,18 +159,11 @@ def __init__( # First convert the relative joint angles to global joint angles in rotation matrix form. if self.smpl_layer.model_type != "flame": - if self.smpl_layer.model_type != "mano": - global_oris = local_to_global( - torch.cat([self.poses_root, self.poses_body, self.poses_left_hand, self.poses_right_hand], dim=-1), - self.skeleton[:, 0], - output_format="rotmat", - ) - else: - global_oris = local_to_global( - torch.cat([self.poses_root, self.poses_body], dim=-1), - self.skeleton[:, 0], - output_format="rotmat", - ) + global_oris = local_to_global( + torch.cat([self.poses_root, self.poses_body], dim=-1), + self.skeleton[:, 0], + output_format="rotmat", + ) global_oris = c2c(global_oris.reshape((self.n_frames, -1, 3, 3))) else: global_oris = np.tile(np.eye(3), self.joints.shape[:-1])[np.newaxis] diff --git a/aitviewer/utils/colors.py b/aitviewer/utils/colors.py new file mode 100644 index 00000000..69bdaf71 --- /dev/null +++ b/aitviewer/utils/colors.py @@ -0,0 +1,43 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import numpy as np + + +def skining_weights_to_color(skinning_weights, alpha): + """Given a skinning weight matrix NvxNj, return a color matrix of shape Nv*3. For each joint Ji i in [0, Nj] , + the color is colors[i]""" + + joints_ids = np.arange(0, skinning_weights.shape[1]) + colors = vertex_colors_from_weights(joints_ids, scale_to_range_1=True, alpha=alpha, shuffle=True, seed=1) + + weights_color = np.matmul(skinning_weights, colors) + return weights_color + + +def vertex_colors_from_weights(weights, scale_to_range_1=True, alpha=None, shuffle=False, seed=0): + """ + Given an array of values of size N, generate an array of colors (Nx3) forming a gradient. + :param weights: Input values (N) + :param scale_to_range_1: If False, the color gradient will cover the values 0 to 1 and plateau beyond + :param alpha: If not None, add an alpha channel of value alpha + :param shuffle: If True, shuffle the colors + :param seed: Seed for the random number generator to shuffle + :return: An array of rgb colors (N, 3). + """ + if scale_to_range_1: + weights = weights - np.min(weights) + weights = weights / np.max(weights) + + from matplotlib import cm + + if alpha is None: + vertex_colors = np.ones((len(weights), 3)) + else: + vertex_colors = alpha * np.ones((len(weights), 4)) + vertex_colors[:, :3] = cm.jet(weights)[:, :3] + + if shuffle: + np.random.seed(seed) + np.random.shuffle(vertex_colors) + + return vertex_colors diff --git a/aitviewer/utils/mocap.py b/aitviewer/utils/mocap.py new file mode 100644 index 00000000..bce2af07 --- /dev/null +++ b/aitviewer/utils/mocap.py @@ -0,0 +1,85 @@ +# # Code Developed by Marilyn Keller, marilyn.keller@tuebingen.mpg.de +# # Do not share or distribute without permission of the author +import numpy as np +import tqdm + +try: + import nimblephysics as nimble +except ImportError: + raise ("nimblephysics not found. Please install nimblephysics with 'pip install nimblephysics' to use this module.") + + +def clean_CMU_mocap_labels(c3dFile: nimble.biomechanics.C3D): + "Rename all the labels with the pattern AAAA-XX and replace them by AAAA" + + c3dFile.markers = [name for name in c3dFile.markers if "-" not in name] + + markerTimesteps = c3dFile.markerTimesteps.copy() + + for markers_dict in markerTimesteps: + markers_dict_clean = markers_dict.copy() + for key in markers_dict: + if "-" in key: + key_clean = key.split("-")[0] + markers_dict_clean[key_clean] = markers_dict_clean.pop(key) + markers_dict.clear() + markers_dict.update(markers_dict_clean) + + c3dFile.markerTimesteps = markerTimesteps + + return c3dFile + + +def load_markers(c3d_path, nb_markers_expected=None): + # Load the marker trajectories + import os + + assert os.path.exists(c3d_path), f"File {c3d_path} not found." + try: + import nimblephysics as nimble + except: + raise ImportError("Please install nimblephysics to load c3d files") + + try: + c3dFile: nimble.biomechanics.C3D = nimble.biomechanics.C3DLoader.loadC3D(os.path.abspath(c3d_path)) + except Exception as e: + print(f"Error loading c3d file {c3d_path}: {e}") + raise e + + c3dFile = clean_CMU_mocap_labels(c3dFile) + + # This c3dFile.markerTimesteps is cryptonite, it keeps doing weird stuff (aka changing values, or you can not edit it), + # it behaves normaly if you make a copy + markers_data_list = c3dFile.markerTimesteps.copy() + + markers_labels = c3dFile.markers + markers_labels.sort() + nb_markers = len(markers_labels) + + if nb_markers_expected is not None: + assert len(markers_labels) == nb_markers_expected, "Expected {} markers, found {}".format( + nb_markers_expected, len(markers_labels) + ) + print(f"Found {nb_markers} markers: {markers_labels}") + + # List of per frame pc array + markers_array = np.zeros((len(markers_data_list), nb_markers, 3)) # FxMx3 + for frame_id, marker_data in (pbar := tqdm.tqdm(enumerate(markers_data_list))): + pbar.set_description("Generating markers point clouds ") + for marker_id, marker_name in enumerate(markers_labels): + if marker_name in marker_data: + marker_pos = marker_data[marker_name] + if np.any(np.abs(marker_pos) > 10e2): + print( + "Warning: marker {} is too far away on frame {}, will be displayed in (0,0,0)".format( + marker_name, frame_id + ) + ) + marker_pos = np.nan * np.zeros((3)) + else: + marker_pos = np.nan * np.zeros((3)) + markers_array[frame_id, marker_id, :] = marker_pos + + fps = c3dFile.framesPerSecond + + return markers_array, markers_labels, fps diff --git a/aitviewer/utils/utils.py b/aitviewer/utils/utils.py index 7742a908..d4bb5c10 100644 --- a/aitviewer/utils/utils.py +++ b/aitviewer/utils/utils.py @@ -298,7 +298,7 @@ def compute_union_of_bounds(nodes): if len(nodes) == 0: return np.zeros((3, 2)) - bounds = np.array([[np.inf, np.NINF], [np.inf, np.NINF], [np.inf, np.NINF]]) + bounds = np.array([[np.inf, -np.inf], [np.inf, -np.inf], [np.inf, -np.inf]]) for n in nodes: child = n.bounds bounds[:, 0] = np.minimum(bounds[:, 0], child[:, 0]) @@ -310,7 +310,7 @@ def compute_union_of_current_bounds(nodes): if len(nodes) == 0: return np.zeros((3, 2)) - bounds = np.array([[np.inf, np.NINF], [np.inf, np.NINF], [np.inf, np.NINF]]) + bounds = np.array([[np.inf, -np.inf], [np.inf, -np.inf], [np.inf, -np.inf]]) for n in nodes: child = n.current_bounds bounds[:, 0] = np.minimum(bounds[:, 0], child[:, 0]) diff --git a/aitviewer/utils/vtp_to_ply.py b/aitviewer/utils/vtp_to_ply.py new file mode 100644 index 00000000..cbf68196 --- /dev/null +++ b/aitviewer/utils/vtp_to_ply.py @@ -0,0 +1,55 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de +import argparse +import os + +import pyvista # required as a .vtp reader + + +def convert_meshes(src_folder, dst_folder): + src = src_folder + if src[-1] != "/": + src += "/" + if dst_folder is None: + target = src + "../Geometry_ply/" + else: + target = dst_folder + if target[-1] != "/": + target += "/" + + os.makedirs(target, exist_ok=True) + + # for each file in src + for filename in os.listdir(src): + ext = os.path.splitext(filename)[-1] + if ext not in [".vtp", ".obj"]: + print("Skipping " + filename) + continue + try: + reader = pyvista.get_reader(src + filename) + except: + print("Could not read " + filename + ". Skipping.") + continue + + mesh = reader.read() + mesh = mesh.triangulate() + # mesh.plot() + + mesh.save(target + filename + ".ply") + print("Converted mesh: " + target + filename + ".ply") + + +if __name__ == "__main__": + # Parse a vtp file and convert it to a ply file + parser = argparse.ArgumentParser(description="Convert a folder of vtp files to a folder of ply files") + parser.add_argument( + "src_folder", + help="folder containing the vtp files to convert", + default="/home/kellerm/Dropbox/MPI/TML/Fullbody_TLModels_v2.0_OS4x/Geometry/", + type=str, + ) + parser.add_argument("dst_folder", help="folder to save the ply files", default=None, type=str) + + args = parser.parse_args() + + src_folder = args.src_folder + dst_folder = args.dst_folder diff --git a/assets/bioamass_screenshot.png b/assets/bioamass_screenshot.png new file mode 100644 index 00000000..3cc1577b Binary files /dev/null and b/assets/bioamass_screenshot.png differ diff --git a/assets/osim_apose.png b/assets/osim_apose.png new file mode 100644 index 00000000..be23a01a Binary files /dev/null and b/assets/osim_apose.png differ diff --git a/assets/skel_sequence.gif b/assets/skel_sequence.gif new file mode 100644 index 00000000..efa56c72 Binary files /dev/null and b/assets/skel_sequence.gif differ diff --git a/examples/load_SKEL.py b/examples/load_SKEL.py new file mode 100644 index 00000000..f5e4167c --- /dev/null +++ b/examples/load_SKEL.py @@ -0,0 +1,69 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import argparse +import os + +import numpy as np +import torch + +from aitviewer.configuration import CONFIG as C +from aitviewer.renderables.skel import SKELSequence +from aitviewer.viewer import Viewer + +try: + from skel.skel_model import SKEL +except Exception as e: + print("Could not import SKEL, make sure you installed the skel repository.") + raise e + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Load a SKEL model and display it.") + parser.add_argument("-s", "--motion_file", type=str, help="Path to a skel motion file", default=None) + parser.add_argument("-z", "--z_up", help="Rotate the mesh 90 deg", action="store_true") + + args = parser.parse_args() + + skel_model = SKEL(gender="female", model_path=C.skel_models).to(C.device) + + if args.motion_file is None: + F = 120 + pose = torch.zeros(F, 46).to(C.device) + betas = torch.zeros(F, 10).to(C.device) + betas[: F // 2, 0] = torch.linspace(-2, 2, F // 2) # Vary beta0 between -2 and 2 + betas[F // 2 :, 1] = torch.linspace(-2, 2, F // 2) # Vary beta1 between -2 and 2 + + trans = torch.zeros(F, 3).to(C.device) + + # Test SKEL forward pass + skel_output = skel_model(pose, betas, trans) + + skel_seq = SKELSequence( + skel_layer=skel_model, + betas=betas, + poses_body=pose, + poses_type="skel", + trans=trans, + is_rigged=True, + show_joint_angles=True, + name="SKEL", + z_up=False, + skin_coloring="skinning_weights", # "pose_offsets", "skinning_weights" + ) + cam_pose = None + + else: + assert os.path.exists( + args.motion_file + ), f"Could not find {args.motion_file}, please provide a valid path to a skel motion file." + assert args.motion_file.endswith(".pkl"), f"Please provide a .pkl file." + + skel_seq = SKELSequence.from_pkl(args.motion_file, name="SKEL", fps_in=120, fps_out=30, z_up=args.z_up) + cam_pose = np.array([0, 1.2, 4.0]) + + v = Viewer() + v.playback_fps = 30 + v.scene.add(skel_seq) + v.run_animations = True + if cam_pose is not None: + v.scene.camera.position = cam_pose + v.run() diff --git a/examples/load_bioamass.py b/examples/load_bioamass.py new file mode 100644 index 00000000..e06faf64 --- /dev/null +++ b/examples/load_bioamass.py @@ -0,0 +1,59 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import os + +import numpy as np + +from aitviewer.configuration import CONFIG as C +from aitviewer.renderables.osim import OSIMSequence +from aitviewer.renderables.smpl import SMPLSequence +from aitviewer.viewer import Viewer + +if __name__ == "__main__": + subj_name = "01" + seq_name = "03" + + c = (149 / 255, 85 / 255, 149 / 255, 0.5) + + to_display = [] + + amass_file = os.path.join(C.datasets.amass, f"CMU/{subj_name}/{subj_name}_{seq_name}_poses.npz") + osim_file = os.path.join(C.datasets.bioamass, f"CMU/{subj_name}/ab_fits/Models/optimized_scale_and_markers.osim") + mot_file = os.path.join(C.datasets.bioamass, f"CMU/{subj_name}/ab_fits/IK/{seq_name}_ik.mot") + + if os.path.exists(C.datasets.amass): + seq_amass = SMPLSequence.from_amass( + npz_data_path=amass_file, + fps_out=30.0, + color=c, + name=f"AMASS {subj_name} {seq_name}", + show_joint_angles=False, + ) + to_display.append(seq_amass) + else: + seq_amass = None + print(f"Could not find AMASS dataset at {C.datasets.amass}. Skipping loading SMPL body.") + + osim_seq = OSIMSequence.from_files( + osim_path=osim_file, + mot_file=mot_file, + name=f"BSM {subj_name} {seq_name}", + fps_out=30, + color_skeleton_per_part=True, + show_joint_angles=False, + is_rigged=False, + ) + + to_display.append(osim_seq) + + # Display in the viewer + v = Viewer() + v.run_animations = True + v.scene.camera.position = np.array([10.0, 2.5, 0.0]) + v.scene.add(*to_display) + + if seq_amass is not None: + v.lock_to_node(seq_amass, (2, 0.7, 2), smooth_sigma=5.0) + v.playback_fps = 30 + + v.run() diff --git a/examples/load_markers.py b/examples/load_markers.py new file mode 100644 index 00000000..d62fec5d --- /dev/null +++ b/examples/load_markers.py @@ -0,0 +1,46 @@ +"""Visualize amass with mocap markers""" + +import os + +import numpy as np + +from aitviewer.configuration import CONFIG as C +from aitviewer.renderables.markers import Markers +from aitviewer.renderables.point_clouds import PointClouds +from aitviewer.renderables.smpl import SMPLSequence +from aitviewer.viewer import Viewer + +if __name__ == "__main__": + # Display in the viewer. + v = Viewer() + + fps_in = 120 # TODO load this from file + fps_out = 60 + + seq_subj = "01" + seq_trial = "03" + + seq_path = f"CMU/{seq_subj}/{seq_subj}_{seq_trial}_poses.npz" + c3d_file = f"CMU/subjects/{seq_subj}/{seq_subj}_{seq_trial}.c3d" + c3d_file_path = os.path.join(C.datasets.amass_mocap, c3d_file) + + c = (85 / 255, 85 / 255, 255 / 255, 1) + markers_pc = Markers.from_c3d(c3d_file_path, color=c, fps_out=fps_out, point_size=15, nb_markers_expected=41) + + # Amass sequence + c = (149 / 255, 85 / 255, 149 / 255, 0.5) + seq_amass = SMPLSequence.from_amass( + npz_data_path=os.path.join(C.datasets.amass, seq_path), + fps_out=fps_out, + color=c, + name="AMASS SMPL sequence", + show_joint_angles=True, + ) + + v.run_animations = True + v.scene.camera.position = np.array([10.0, 2.5, 0.0]) + + v.scene.add(seq_amass) + v.scene.add(markers_pc) + + v.run() diff --git a/examples/load_osim.py b/examples/load_osim.py new file mode 100644 index 00000000..dd07057d --- /dev/null +++ b/examples/load_osim.py @@ -0,0 +1,130 @@ +# Copyright (C) 2024 Max Planck Institute for Intelligent Systems, Marilyn Keller, marilyn.keller@tuebingen.mpg.de + +import argparse +import os +from typing import Optional + +import numpy as np + +from aitviewer.configuration import CONFIG as C +from aitviewer.renderables.markers import Markers +from aitviewer.renderables.osim import OSIMSequence +from aitviewer.utils.vtp_to_ply import convert_meshes +from aitviewer.viewer import Viewer + + +def display_model_in_viewer( + osim: Optional[str] = None, + mot: Optional[str] = None, + fps: Optional[int] = None, + color_parts: bool = False, + color_markers: bool = False, + mocap: Optional[str] = None, + joints: bool = False, +): + """ + Load an OpenSim model and a motion file and display it in the viewer. + Args: + osim (str, optional): Path to the osim file. If no file is specified, the default Rajagopal gait model will be loaded from nimble. Defaults to None. + mot (str, optional): Path to the motion file. Defaults to None. + fps (int, optional): Generating the meshes for all the frames can take a long time and a lot of memory. Use a low fps to avoid this problem. Defaults to 30. + color_parts (bool, optional): Color the skeleton by parts, as defined in the .osim file. Defaults to False. + color_markers (bool, optional): Each marker is attached to a skeleton part. This option colors the markers to match the parent part color. Defaults to False. + mocap (str, optional): If a Mocap file is specified, display the markers mocap with their labels. For now, only the .c3d format is supported. Defaults to None. + joints (bool, optional): Show model joints as spheres. Defaults to False. + """ + + if args.osim is not None: + # Check that a folder named Geometry exists in the same folder as the osim file. + osim_dir = os.path.dirname(args.osim) + geom_dir = os.path.join(osim_dir, "Geometry") + if not os.path.exists(geom_dir): + print( + "Geometry folder does not exist in the same folder as the osim file. Please add a Geometry folder containing the skeleton .vtp, .obj or .ply file in the same folder as the provided .osim. " + ) + exit(1) + + # Check that the Geometry folder contains at least a file of type .ply + ply_files = [f for f in os.listdir(geom_dir) if f.endswith(".ply")] + if len(ply_files) == 0: + print("Geometry folder does not contain any .ply files.") + print("The provided folder meshes will be converted to .ply. Press a key to continue or CTRL-C to abort.") + + # Convert the provided meshes to .ply + convert_meshes(geom_dir, geom_dir) + + if args.mot is None: + osim_seq = OSIMSequence.a_pose( + args.osim, + name="OpenSim template", + show_joint_angles=args.joints, + color_skeleton_per_part=args.color_parts, + color_markers_per_part=args.color_markers, + ) + + else: + osim_seq = OSIMSequence.from_files( + args.osim, + args.mot, + show_joint_angles=args.joints, + color_skeleton_per_part=args.color_parts, + color_markers_per_part=args.color_markers, + fps_out=args.fps, + ) + + v = Viewer() + v.scene.add(osim_seq) + + v.lock_to_node(osim_seq, (5, 2, 0), smooth_sigma=5.0) + + if args.mocap is not None: + # check that the mocap file is in .c3d format + assert args.mocap.endswith(".c3d"), "Mocap file must be in .c3d format." + marker_seq = Markers.from_c3d(args.mocap, fps_out=args.fps, color=[0, 255, 0, 255]) + v.scene.add(marker_seq) + + if args.fps is not None: + v.playback_fps = args.fps + + v.run_animations = True + + v.run() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Load an OpenSim model and a motion file and display it in the viewer." + ) + + parser.add_argument( + "--osim", + type=str, + help="Path to the osim file. If no file is specified, the default Rajagopal gait model will be loaded from nimble", + default=None, + ) + parser.add_argument("--mot", type=str, help="Path to the motion file", default=None) + parser.add_argument( + "--fps", + type=int, + default=None, + help="Generating the meshes for all the frames can take a long time and a lot of memory. Use a low fps to avoid this problem.", + ) + + parser.add_argument( + "--color_parts", action="store_true", help="Color the skeleton by parts, as defined in the .osim file." + ) + parser.add_argument( + "--color_markers", + action="store_true", + help="Each marker is attached to a skeleton part. This option colors the markers to match the parent part color.", + ) + parser.add_argument( + "--mocap", + type=str, + help="If a Mocap file is specified, display the markers mocap with their labels. For now, only the .c3d format is supported.", + ) + parser.add_argument("--joints", action="store_true", help="Show model joints as spheres.") + + args = parser.parse_args() + + display_model_in_viewer(**args.__dict__) diff --git a/setup.py b/setup.py index afb07dc8..6368cc7b 100644 --- a/setup.py +++ b/setup.py @@ -39,9 +39,9 @@ description="Viewing and rendering of sequences of 3D data.", long_description=open("README.md").read(), long_description_content_type="text/markdown", - url="https://github.com/eth-ait/aitviewer", + url="https://github.com/eth-ait/aitviewer-skel", version=__version__, - author="Manuel Kaufmann, Velko Vechev, Dario Mylonopoulos", + author="Manuel Kaufmann, Velko Vechev, Dario Mylonopoulos, Marilyn Keller", packages=find_packages(), include_package_data=True, keywords=[ @@ -57,11 +57,13 @@ "visualization", ], platforms=["any"], - python_requires=">=3.7,<3.11", + + python_requires=">=3.7", + install_requires=requirements, project_urls={ "Documentation": "https://eth-ait.github.io/aitviewer/", - "Source": "https://github.com/eth-ait/aitviewer", + "Source": "https://github.com/eth-ait/aitviewer-skel", "Bug Tracker": "https://github.com/eth-ait/aitviewer/issues", }, )