trc from mot and osim files

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'<input_mot_file>', r'<output_osim_file>', r'<output_trc_file>')
+    python -m trc_from_mot_osim -m "<input_mot_file>" -o "<input_osim_file>"
+    python -m trc_from_mot_osim -m "<input_mot_file>" -o "<input_osim_file>" -t "<output_trc_file>"
+'''
+
+
+## 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__ = "[email protected]"
+__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'<input_mot_file>', r'<output_osim_file>', r'<trc_output_file>')
+    python -m trc_from_mot_osim -m "<input_mot_file>" -o "<input_osim_file>"
+    python -m trc_from_mot_osim -m "<input_mot_file>" -o "<input_osim_file>" -t "<trc_output_file>"
+    '''
+
+    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)

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.
    </pre>
 </details>