Spaces:
Starting
on
L4
Starting
on
L4
File size: 4,566 Bytes
5ac1897 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
# 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)
@classmethod
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
@classmethod
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
@classmethod
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
@classmethod
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
|