Spaces:
Sleeping
Sleeping
| # Copyright (C) 2024 Max Planck Institute for Intelligent Systems Tuebingen, Marilyn Keller | |
| import os | |
| import tqdm | |
| import numpy as np | |
| import yaml | |
| import nimblephysics as nimble | |
| from skel import kin_skel | |
| class OSIM: | |
| def __init__(self, osim_path, motion_path): | |
| self.osim_path = osim_path | |
| self.motion_path = motion_path | |
| self.osim = OSIM.load_osim(osim_path) | |
| self.motion = OSIM.load_motion(self.osim, motion_path) | |
| self.n_frames = self.motion.shape[0] | |
| self.markers_labels = OSIM.get_markers_labels(self.osim) | |
| self.joints_labels = OSIM.get_joints_labels(self.osim) | |
| def load_osim(cls, osim_path): | |
| """Load an osim file""" | |
| assert os.path.exists(osim_path), "Could not find osim file: {}".format(osim_path) | |
| 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 | |
| def load_motion(cls, osim, motion_path): | |
| """Load a .mot motion file""" | |
| assert os.path.exists(motion_path), "Could not find motion file: {}".format(motion_path) | |
| mot: nimble.biomechanics.OpenSimMot = nimble.biomechanics.OpenSimParser.loadMot(osim.skeleton, motion_path) | |
| assert mot is not None, "Could not load motion file: {}".format(motion_path) | |
| motion = np.array(mot.poses.T) | |
| return motion | |
| def get_markers_labels(cls, osim): | |
| """Return a list of markers in the osim file""" | |
| markers_labels = [ml for ml in osim.markersMap.keys()] | |
| # markers_labels.sort() | |
| return markers_labels | |
| def get_joints_labels(cls, osim): | |
| node_names = [n.getName() for n in osim.skeleton.getBodyNodes()] | |
| return node_names | |
| def skel_joint_mask(self, skel_mapping): | |
| """ Return a mask to map the joints of the skel model to the joints of the osim model""" | |
| corresp_skel_joint_ids = [] # For each osim joint, the index of the corresponding skel joint | |
| joint_mapping_dict = skel_mapping['joints_mapping'] | |
| skel_joint_names = kin_skel.skel_joints_name | |
| for osim_joint_name in joint_mapping_dict.keys(): | |
| skel_joint_name = joint_mapping_dict[osim_joint_name] | |
| skel_joint_index = skel_joint_names.index(skel_joint_name) | |
| corresp_skel_joint_ids.append(skel_joint_index) | |
| import ipdb; ipdb.set_trace() | |
| return corresp_skel_joint_ids | |
| def skel_marker_mask(self, skel_mapping): | |
| markers_mapping_dict = skel_mapping['markers_mapping'] | |
| osim_mask = [] | |
| skel_verts_mask = [] | |
| for osim_marker_name in self.markers_labels: | |
| osim_marker_index = self.markers_labels.index(osim_marker_name) | |
| osim_mask.append(osim_marker_index) | |
| skel_verts_mask.append(markers_mapping_dict[osim_marker_name]) | |
| return osim_mask, skel_verts_mask | |
| def run_fk(self): | |
| """Return a list of markers in the osim file""" | |
| markers = np.zeros((self.n_frames, len(self.markers_labels), 3)) | |
| joints = np.zeros([self.n_frames, len(self.joints_labels), 3]) | |
| joints_ori = np.zeros([self.n_frames, len(self.joints_labels), 3, 3]) | |
| for frame_id in (pbar := tqdm.tqdm(range(self.n_frames))): | |
| pbar.set_description("Running OpenSim forward kinematics") | |
| pose = self.motion[frame_id, :] | |
| # 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.joints_labels): | |
| # part transfo | |
| transfo = self.osim.skeleton.getBodyNode(node_name).getWorldTransform() | |
| # Update joint | |
| joints[frame_id, ni, :] = transfo.translation() | |
| joints_ori[frame_id, ni, :, :] = transfo.rotation() | |
| return markers, joints, joints_ori | |