diff --git a/Pose2Sim/Utilities/trc_from_mot_osim.py b/Pose2Sim/Utilities/trc_from_mot_osim.py new file mode 100644 index 0000000..ffcf787 --- /dev/null +++ b/Pose2Sim/Utilities/trc_from_mot_osim.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + + +''' + ################################################## + ## Build trc from mot and osim files ## + ################################################## + + Build a trc file which stores all real and virtual markers + calculated from a .mot motion file and a .osim model file. + + Usage: + from Pose2Sim.Utilities import trc_from_mot_osim; trc_from_mot_osim.trc_from_mot_osim_func(r'', r'', r'') + python -m trc_from_mot_osim -m "" -o "" + python -m trc_from_mot_osim -m "" -o "" -t "" +''' + + +## INIT +import os +import pandas as pd +import numpy as np +import opensim as osim +import argparse + + +## AUTHORSHIP INFORMATION +__author__ = "David Pagnon" +__copyright__ = "Copyright 2021, Pose2Sim" +__credits__ = ["David Pagnon"] +__license__ = "BSD 3-Clause License" +__version__ = "0.4" +__maintainer__ = "David Pagnon" +__email__ = "contact@david-pagnon.com" +__status__ = "Development" + + +## FUNCTIONS +def get_marker_positions(motion_data, model): + ''' + Get dataframe of marker positions + + INPUTS: + - motion_data: .mot file opened with osim.TimeSeriesTable + - model: .osim file opened with osim.Model + + OUTPUT: + - marker_positions_pd: DataFrame of marker positions + ''' + + # Markerset + marker_set = model.getMarkerSet() + marker_set_names = [mk.getName() for mk in list(marker_set)] + marker_set_names_xyz = np.array([[m+'_x', m+'_y', m+'_z'] for m in marker_set_names]).flatten() + + # Data + times = motion_data.getIndependentColumn() + joint_angle_set_names = motion_data.getColumnLabels() # or [c.getName() for c in model.getCoordinateSet()] + motion_data_pd = pd.DataFrame(motion_data.getMatrix().to_numpy(), columns=joint_angle_set_names) + + # Get marker positions at each state + state = model.initSystem() + marker_positions = [] + for n,t in enumerate(times): + # put the model in the right position + for coord in joint_angle_set_names: + if not coord.endswith('_tx') and not coord.endswith('_ty') and not coord.endswith('_tz'): + value = motion_data_pd.loc[n,coord]*np.pi/180 + else: + value = motion_data_pd.loc[n,coord] + model.getCoordinateSet().get(coord).setValue(state,value) + # get marker positions + marker_positions += [np.array([marker_set.get(mk_name).findLocationInFrame(state, model.getGround()).to_numpy() for mk_name in marker_set_names]).flatten()] + marker_positions_pd = pd.DataFrame(marker_positions, columns=marker_set_names_xyz) + marker_positions_pd.insert(0, 'time', times) + marker_positions_pd.insert(0, 'frame', np.arange(len(times))) + + return marker_positions_pd + + +def trc_from_mot_osim_func(*args): + ''' + Build a trc file which stores all real and virtual markers + calculated from a .mot motion file and a .osim model file. + + Usage: + from Pose2Sim.Utilities import trc_from_mot_osim; trc_from_mot_osim.trc_from_mot_osim_func(r'', r'', r'') + python -m trc_from_mot_osim -m "" -o "" + python -m trc_from_mot_osim -m "" -o "" -t "" + ''' + + try: + motion_path = args[0]['input_mot_file'] # invoked with argparse + osim_path = args[0]['input_osim_file'] + if args[0]['trc_output_file'] == None: + trc_path = motion_path.replace('.mot', '.trc') + else: + trc_path = args[0]['trc_output_file'] + except: + motion_path = args[0] # invoked as a function + osim_path = args[1] + try: + trc_path = args[2] + except: + trc_path = motion_path.replace('.mot', '.trc') + + # Create dataframe with marker positions + model = osim.Model(osim_path) + motion_data = osim.TimeSeriesTable(motion_path) + marker_positions_pd = get_marker_positions(motion_data, model) + + # Trc header + times = motion_data.getIndependentColumn() + marker_set = model.getMarkerSet() + marker_set_names = [mk.getName() for mk in list(marker_set)] + + fps = str(int( 1/ (times[1]-times[0]) / (len(times)-1) )) + nb_frames = str(len(times)) + nb_markers = str(len(marker_set_names)) + header0_str = 'PathFileType\t4\t(X/Y/Z)\t' + os.path.basename(trc_path) + + header1 = {} + header1['DataRate'] = fps + header1['CameraRate'] = fps + header1['NumFrames'] = nb_frames + header1['NumMarkers'] = nb_markers + header1['Units'] = 'm' + header1['OrigDataRate'] = fps + header1['OrigDataStartFrame'] = '0' + header1['OrigNumFrames'] = nb_frames + header1_str1 = '\t'.join(header1.keys()) + header1_str2 = '\t'.join(header1.values()) + + header2_str1 = 'Frame#\tTime\t' + '\t\t\t'.join([mk.strip() for mk in marker_set_names]) + '\t\t' + header2_str2 = '\t\t'+'\t'.join(['X{i}\tY{i}\tZ{i}'.format(i=i+1) for i in range(int(header1['NumMarkers']))]) + + header_trc = '\n'.join([header0_str, header1_str1, header1_str2, header2_str1, header2_str2]) + + # write data + with open(trc_path, 'w') as trc_o: + trc_o.write(header_trc+'\n') + marker_positions_pd.to_csv(trc_path, header=False, sep = '\t', mode='a', index=False) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('-m', '--input_mot_file', required = True, help='input mot file') + parser.add_argument('-o', '--input_osim_file', required = True, help='input osim file') + parser.add_argument('-t', '--trc_output_file', required=False, help='trc output file') + args = vars(parser.parse_args()) + + trc_from_mot_osim_func(args) \ No newline at end of file diff --git a/README.md b/README.md index 4a8e016..1ebd161 100644 --- a/README.md +++ b/README.md @@ -730,6 +730,9 @@ Detects gait events from point coordinates according to [Zeni et al. (2008)](htt [trc_combine.py](https://github.com/perfanalytics/pose2sim/blob/main/Pose2Sim/Utilities/trc_combine.py) Combine two trc files, for example a triangulated DeepLabCut trc file and a triangulated OpenPose trc file. + +[trc_from_mot_osim.py](https://github.com/perfanalytics/pose2sim/blob/main/Pose2Sim/Utilities/trc_from_mot_osim.py) +Build a trc file from a .mot motion file and a .osim model file.