kinect_joint_angles/functions.py at main · bs97/kinect_joint_angles · GitHub

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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
import opensim as osim
import pandas as pd
import numpy as np
from scipy.spatial.transform import Rotation as R


# This function uses the opensim analyze tool to create a states file from a coordinates file
def create_states_file_from_coordinates_file(analyze_settings_template_file, model_file, coord_file, results_path,
                                             start_time, end_time):
    # Instantiate a Analyze Tool
    analyze_tool = osim.AnalyzeTool(analyze_settings_template_file)
    analyze_tool.setModelFilename(model_file)
    analyze_tool.setResultsDir(results_path)
    analyze_tool.setCoordinatesFileName(coord_file)
    analyze_tool.setName("OMC")
    analyze_tool.run()


def run_analyze_tool(analyze_settings_template_file, results_dir, model_file_path, mot_file_path, start_time, end_time):
    analyze_tool = osim.AnalyzeTool(analyze_settings_template_file)
    analyze_tool.updAnalysisSet().cloneAndAppend(osim.BodyKinematics())
    analyze_tool.setModelFilename(model_file_path)
    analyze_tool.setName("analyze")
    analyze_tool.setCoordinatesFileName(mot_file_path)
    analyze_tool.setStartTime(start_time)
    analyze_tool.setFinalTime(end_time)
    analyze_tool.setResultsDir(results_dir)
    print('Running Analyze Tool...')
    analyze_tool.run()
    print('Analyze Tool run finished.')


def get_body_quats_from_analysis_sto(analysis_sto_path, start_time, end_time):
    analysis_table = osim.TimeSeriesTable(analysis_sto_path)
    analysis_table.trim(start_time, end_time)

    column_labels = analysis_table.getColumnLabels()
    print("Available columns in the analysis .sto file:")
    for label in column_labels:
        print(label)

    # Update these lines to match the actual column names in your .sto file
    thorax_Ox = analysis_table.getDependentColumn('thorax_Ox').to_numpy()
    thorax_Oy = analysis_table.getDependentColumn('thorax_Oy').to_numpy()
    thorax_Oz = analysis_table.getDependentColumn('thorax_Oz').to_numpy()
    humerus_Ox = analysis_table.getDependentColumn('humerus_r_Ox').to_numpy()
    humerus_Oy = analysis_table.getDependentColumn('humerus_r_Oy').to_numpy()
    humerus_Oz = analysis_table.getDependentColumn('humerus_r_Oz').to_numpy()
    radius_Ox = analysis_table.getDependentColumn('radius_r_Ox').to_numpy()
    radius_Oy = analysis_table.getDependentColumn('radius_r_Oy').to_numpy()
    radius_Oz = analysis_table.getDependentColumn('radius_r_Oz').to_numpy()

    thorax_eulers = np.stack((thorax_Ox, thorax_Oy, thorax_Oz), axis=1)
    humerus_eulers = np.stack((humerus_Ox, humerus_Oy, humerus_Oz), axis=1)
    radius_eulers = np.stack((radius_Ox, radius_Oy, radius_Oz), axis=1)

    # Create an array of scipy Rotations
    thorax_R = R.from_euler('XYZ', thorax_eulers, degrees=True)
    humerus_R = R.from_euler('XYZ', humerus_eulers, degrees=True)
    radius_R = R.from_euler('XYZ', radius_eulers, degrees=True)

    thorax_quats = thorax_R.as_quat()
    humerus_quats = humerus_R.as_quat()
    radius_quats = radius_R.as_quat()

    return thorax_quats, humerus_quats, radius_quats


def get_body_quats_from_states(states_file, model_file, results_path, analyze_settings_template_file, mot_file_path,
                               start_time, end_time):
    # Run Analyze Tool to generate the .sto file
    run_analyze_tool(analyze_settings_template_file, results_path, model_file, mot_file_path, start_time, end_time)

    # Path to the generated analysis .sto file
    analysis_sto_path = f'{results_path}/analyze_BodyKinematics_pos_global.sto'

    # Extract the body quaternions from the analysis .sto file
    thorax_quats, humerus_quats, radius_quats = get_body_quats_from_analysis_sto(analysis_sto_path, start_time,
                                                                                 end_time)

    thorax_df = pd.DataFrame(thorax_quats, columns=['w', 'x', 'y', 'z'])
    humerus_df = pd.DataFrame(humerus_quats, columns=['w', 'x', 'y', 'z'])
    radius_df = pd.DataFrame(radius_quats, columns=['w', 'x', 'y', 'z'])

    thorax_df.to_csv(f'{results_path}/thorax_orientations.csv', index=False)
    humerus_df.to_csv(f'{results_path}/humerus_orientations.csv', index=False)
    radius_df.to_csv(f'{results_path}/radius_orientations.csv', index=False)

    return True


# Function to read quaternions from the CSV files
def read_in_quats(start_time, end_time, file_name, trim_bool):
    thorax_df = pd.read_csv(file_name.replace('Body_Oris.csv', 'thorax_orientations.csv'))
    humerus_df = pd.read_csv(file_name.replace('Body_Oris.csv', 'humerus_orientations.csv'))
    radius_df = pd.read_csv(file_name.replace('Body_Oris.csv', 'radius_orientations.csv'))

    if trim_bool:
        thorax_df = thorax_df[(thorax_df['time'] >= start_time) & (thorax_df['time'] <= end_time)]
        humerus_df = humerus_df[(humerus_df['time'] >= start_time) & (humerus_df['time'] <= end_time)]
        radius_df = radius_df[(radius_df['time'] >= start_time) & (radius_df['time'] <= end_time)]

    thorax_quats = thorax_df.to_numpy()
    humerus_quats = humerus_df.to_numpy()
    radius_quats = radius_df.to_numpy()

    return thorax_quats, humerus_quats, radius_quats


# Define a function for quaternion multiplication
def quat_mul(Q0, Q1):
    """
    Multiplies two quaternions.
    Input
    :param Q0: A 4 element array containing the first quaternion (q01,q11,q21,q31)
    :param Q1: A 4 element array containing the second quaternion (q02,q12,q22,q32)
    Output
    :return: A 4 element array containing the final quaternion (q03,q13,q23,q33)
    """
    w0 = Q0[0]
    x0 = Q0[1]
    y0 = Q0[2]
    z0 = Q0[3]
    w1 = Q1[0]
    x1 = Q1[1]
    y1 = Q1[2]
    z1 = Q1[3]
    # Computer the product of the two quaternions, term by term
    Q0Q1_w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
    Q0Q1_x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
    Q0Q1_y = w0 * y1 - x0 * z1 + y0 * w1 + z0 * x1
    Q0Q1_z = w0 * z1 + x0 * y1 - y0 * x1 + z0 * w1
    # Create a 4 element array containing the final quaternion
    final_quaternion = np.array([Q0Q1_w, Q0Q1_x, Q0Q1_y, Q0Q1_z])
    return final_quaternion


# Calculate quaternion conjugate
def quat_conj(Q0):
    w0 = Q0[0]
    x0 = Q0[1]
    y0 = Q0[2]
    z0 = Q0[3]
    output_quaternion = np.array([w0, -x0, -y0, -z0])
    return output_quaternion


# This function takes quaternion body orientations and outputs the joint euler angles
def get_JA_euls_from_quats(body1_quats, body2_quats, eul_seq):

    n_rows = len(body1_quats)
    eul_1_arr = np.zeros((n_rows))
    eul_2_arr = np.zeros((n_rows))
    eul_3_arr = np.zeros((n_rows))
    for row in range(n_rows):
        joint_Rot = quat_mul(quat_conj(body1_quats[row]), body2_quats[row])  # Calculate joint Rot quat
        joint_scipyR = R.from_quat([joint_Rot[1], joint_Rot[2], joint_Rot[3], joint_Rot[0]])    # In scalar last format
        joint_eul = joint_scipyR.as_euler(eul_seq, degrees=True)    # Get euler angles
        eul_1_arr[row], eul_2_arr[row], eul_3_arr[row] = joint_eul[0], joint_eul[1], joint_eul[2]


    return eul_1_arr, eul_2_arr, eul_3_arr