first submission

Author: Tian-A-Tan

README.md

@@ -1 +1,29 @@
-# MotionDiffusionModel
+
+# GaitDynamics: A Foundation Model for Analyzing Gait Dynamics 
+By Tian Tan, Tom Van Wouwe, Keenon F. Werling, Scott L. Delp, Jennifer L. Hicks, C. Karen Liu, and Akshay S. Chaudhari
+
+## Exclusive Summary
+GaitDynamics is a generative foundation model for general-purpose gait dynamics prediction.
+We illustrate in three diverse tasks with different inputs, outputs, and clinical impacts: i) estimating 
+external forces from kinematics, ii) predicting the influence of gait modifications on knee loading without human 
+experiments, and iii) predicting comprehensive kinematics and kinetic changes that occur with increasing running 
+speeds.
+
+## Corresponding Publication
+This repository includes the code and models for an [abstract](./figures/readme_fig/Tan_ASB2024.pdf).
+Full-length preprint is coming soon.
+
+## Environment
+Our code is developed under the following environment. Versions different from ours may still work.
+
+Python 3.9.16; Pytorch 1.13.1; Cuda 11.6; Cudnn 8.3.2; numpy 1.23.5;
+
+## Dataset
+[AddBiomechanics Dataset](https://addbiomechanics.org/download_data.html)
+
+## Example code
+[A Google Colab notebook](https://colab.research.google.com/drive/1n6kH3gnwLdQ2DH5krigbkiO06NjDtyxI?usp=sharing)
+is provided for the downstream tasks 1 – force estimation using flexible combinations of kinematic inputs.
+By executing the code, and example .mot file with joint angles of an OpenSim skeletal model will be imported from GitHub.
+Users can upload their own .mot files to the Colab notebook to obtain force predictions.
+To use a reduced kinematic input combinations, simply delete the corresponding columns in the .mot file.

data/addb_dataset.py

@@ -40,6 +40,7 @@ def __init__(
             restrict_contact_bodies=True,
             use_camargo_lumbar_reconstructed=False,
             check_cop_to_calcn_distance=True,
+            wrong_cop_ratio=0.002
     ):
         self.data_path = data_path
         self.trial_start_num = trial_start_num
@@ -54,6 +55,7 @@ def __init__(
         self.restrict_contact_bodies = restrict_contact_bodies
         self.use_camargo_lumbar_reconstructed = use_camargo_lumbar_reconstructed
         self.check_cop_to_calcn_distance = check_cop_to_calcn_distance
+        self.wrong_cop_ratio = wrong_cop_ratio
         self.skels = {}
         self.num_of_excluded_trials = {'contact_body_num': 0, 'trial_length': 0, 'lumbar_rotation': 0, 'wrong_cop': 0,
                                        'large_moving_direction_change': 0, 'jittery_sample': 0}
@@ -513,7 +515,7 @@ def load_addb(self, opt, max_trial_num):
                     #       f' flipped {len(grf_flag_counts)} times, thus setting all to True.', end='')
                     probably_missing = [False] * len(probably_missing)
 
-                states = norm_cops(skel, states, opt, weight_kg, height_m, self.check_cop_to_calcn_distance)
+                states = norm_cops(skel, states, opt, weight_kg, height_m, self.check_cop_to_calcn_distance, self.wrong_cop_ratio)
                 if states is False:
                     print(f'{sub_and_trial_name} has CoP far away from foot, skipping')
                     self.num_of_excluded_trials['wrong_cop'] += 1

data/check_b3d.py

@@ -0,0 +1,23 @@
+import nimblephysics as nimble
+
+
+b3d_file = '/mnt/d/Downloads/AB06_split0.b3d'
+subject = nimble.biomechanics.SubjectOnDisk(b3d_file)
+for trial_id in range(subject.getNumTrials()):
+    trial_name = subject.getTrialName(trial_id)
+    print(trial_name)
+
+    missing_grf_labels = subject.getMissingGRF(trial_id)
+
+    probably_missing = [reason != nimble.biomechanics.MissingGRFReason.notMissingGRF for reason
+                        in subject.getMissingGRF(trial_id)]
+
+
+
+
+
+
+
+
+
+

example_usage/compress_model.py

@@ -133,7 +133,7 @@ def compress_model(checkpoint_path, model_class, new_model_name):
     compress_model(checkpoint_path, DanceDecoder, 'GaitDynamicsDiffusion')
 
     # # full-body tf
-    # checkpoint_path = os.getcwd() + '/../trained_models/train-2560_tf.pt'
+    # checkpoint_path = os.getcwd() + '/../trained_models/train-7680_tf.pt'
     # compress_model(checkpoint_path, TransformerEncoderArchitecture, 'GaitDynamicsRefinement')
 
     # # hip-knee tf

example_usage/gait_dynamics.py

@@ -1208,6 +1208,7 @@ def convertDfToGRFMot(df, out_folder, dt, max_time=None):
             out_file.write('\t' + str(0))
         out_file.write('\n')
     out_file.close()
+    print('GRF file exported to ' + out_folder)
 
 
 def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:

figures/da_grf_test_set_0.py

@@ -72,9 +72,9 @@ def load_baseline_model(opt, model_to_test):
     elif model_to_test == 3:
         model_architecture_class = SugaiNetArchitecture
         if opt.use_server:
-            opt.checkpoint_bl = opt.data_path_parent + f"/../code/runs/train/{'SugaiNetArchitecture_ema999'}/weights/train-{'2000_sugainet'}.pt"
+            opt.checkpoint_bl = opt.data_path_parent + f"/../code/runs/train/{'SugaiNetArchitecture_ema999'}/weights/train-{'2020_sugainet'}.pt"
         else:
-            opt.checkpoint_bl = os.path.dirname(os.path.realpath(__file__)) + f"/../trained_models/train-{'2000_sugainet'}.pt"
+            opt.checkpoint_bl = os.path.dirname(os.path.realpath(__file__)) + f"/../trained_models/train-{'2020_sugainet'}.pt"
         model_key = 'sugainet'
 
     set_with_arm_opt(opt, False)
@@ -403,6 +403,10 @@ def load_test_dataset_dict():
     for dset in DATASETS_NO_ARM:
         if dset in dset_to_skip:
             continue
+        if 'vanderZee2022' in dset:
+            wrong_cop_ratio = 0.01
+        else:
+            wrong_cop_ratio = 0.002
         print(dset)
         test_dataset = MotionDataset(
             data_path=opt.data_path_test,
@@ -414,6 +418,7 @@ def load_test_dataset_dict():
             include_trials_shorter_than_window_len=True,
             restrict_contact_bodies=False,
             max_trial_num=max_trial_num,
+            wrong_cop_ratio=wrong_cop_ratio
         )
         test_dataset_dict[dset] = test_dataset
     return test_dataset_dict

figures/da_grf_test_set_1.py

@@ -7,7 +7,7 @@
 from da_grf_test_set_0 import cols_to_unmask, dset_to_skip, drop_frame_num_range
 from data.addb_dataset import MotionDataset
 from matplotlib import rc, lines
-from fig_utils import FONT_DICT_SMALL, FONT_SIZE_SMALL, format_axis, LINE_WIDTH
+from fig_utils import FONT_DICT_SMALL, FONT_SIZE_SMALL, format_axis, LINE_WIDTH, FONT_DICT_X_SMALL
 from scipy.stats import friedmanchisquare, wilcoxon
 
 
@@ -248,7 +248,7 @@ def get_all_the_metrics(model_key):
         #     plt.plot(true_concat[within_gait_cycle, param_col_loc])
         #     plt.plot(pred_concat[within_gait_cycle, param_col_loc])
         #     plt.title(dset_short + ' ' + str(metric_dset[param]))
-        #     plt.show()
+        # plt.show()
 
     # for param_col, metric_list in metric_all_dsets.items():
     #     if param_col == 'dset_short':
@@ -259,13 +259,18 @@ def get_all_the_metrics(model_key):
 
 def draw_fig_2(fast_run=False):
     def format_ticks(ax):
-        ax.set_ylabel('Mean Absolute Error (\% Body Weight)', fontdict=FONT_DICT_SMALL)
+        ax.text(-0.7, 5, 'Mean Absolute Error (\% BW)', rotation=90, fontdict=FONT_DICT_SMALL, verticalalignment='center')
+        ax.text(-0.8, 12.5, 'Better', rotation=90, fontdict=FONT_DICT_SMALL, color='green', verticalalignment='center')
+        ax.annotate('', xy=(-0.08, 0.78), xycoords='axes fraction', xytext=(-0.08, 0.98),
+                    arrowprops=dict(arrowstyle="->", color='green'))
+
         ax.set_yticks(range(0, 15, 2))
         ax.set_yticklabels(range(0, 15, 2), fontdict=FONT_DICT_SMALL)
-        # ax.set_xticks([0.25, 1.25, 2.25, 3.25])
+        ax.set_xlim([-0.3, 3.8])
+
         ax.set_xticks([])
         for i in range(4):
-            ax.text(i+0.25, -1.1, list(params_name_formal_name_pairs.values())[i], fontdict=FONT_DICT_SMALL, ha='center')
+            ax.text(i+0.25, -1.6, list(params_name_formal_name_pairs.values())[i], fontdict=FONT_DICT_SMALL, ha='center')
 
     colors = [np.array(x) / 255 for x in [[70, 130, 180], [207, 154, 130], [177, 124, 90]]]        #  [207, 154, 130], [100, 155, 227]
     folder = 'fast' if fast_run else 'full'
@@ -274,14 +279,16 @@ def format_ticks(ax):
     metric_sugainet = get_all_the_metrics(model_key=f'/{folder}/sugainet_none_diffusion_filling')
 
     params_name_formal_name_pairs = {
-        'calcn_l_force_vy_max': r'$f_v$ - Peak', 'calcn_l_force_vy': r'$f_v$ - Profile',
-        'calcn_l_force_vx': r'$f_{ap}$ - Profile', 'calcn_l_force_vz': r'$f_{ml}$ - Profile'}
+        'calcn_l_force_vy': 'Vertical\nForce (Profile)',
+        'calcn_l_force_vx': 'Anterior-Posterior\nForce (Profile)',
+        'calcn_l_force_vz': 'Medial-Lateral\nForce (Profile)',
+        'calcn_l_force_vy_max': 'Vertical\nForce (Peak)'}
     params_of_interest = list(params_name_formal_name_pairs.keys())
 
     rc('text', usetex=True)
     plt.rc('font', family='Helvetica')
 
-    fig = plt.figure(figsize=(5, 3.5))
+    fig = plt.figure(figsize=(7.7, 4.5))
     print('Parameter\t\tAll\t\t1-2\t\t1-3\t\t2-3')
     for i_axis, param in enumerate(params_of_interest):
         bar_locs = [i_axis, i_axis + 0.25, i_axis + 0.5]
@@ -301,25 +308,35 @@ def format_ticks(ax):
         print()
 
     # From "Comparison of different machine learning models to enhance sacral acceleration-based estimations of running stride temporal variables and peak vertical ground reaction force"
-    line0, = plt.plot([-0.2, 0.7], [13, 13], ':', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    line0, = plt.plot([2.8, 3.7], [13, 13], '--', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    plt.text(4.1, 14, 'Running MDC - Healthy [25]', fontdict=FONT_DICT_SMALL, color=[0.0, 0.0, 0.0], va='center')
+    plt.annotate('', xytext=(4.05, 14), xycoords='data', xy=(3.75, 13), arrowprops=dict(arrowstyle="->"))
+
     # From "Intra-rater repeatability of gait parameters in healthy adults during self-paced treadmill-based virtual reality walking"
-    line1, = plt.plot([-0.2, 0.7], [10.18, 10.18], '--', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    line1, = plt.plot([2.8, 3.7], [10.18, 10.18], '--', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    plt.text(4.1, 11.18, 'Walking MDC - Healthy [26]', fontdict=FONT_DICT_SMALL, color=[0.0, 0.0, 0.0], va='center')
+    plt.annotate('', xytext=(4.05, 11.18), xycoords='data', xy=(3.75, 10.18), arrowprops=dict(arrowstyle="->"))
+
     # and "Minimal detectable change for gait variables collected during treadmill walking in individuals post-stroke"
-    line2, = plt.plot([-0.2, 0.7], [4.65, 4.65], '-', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    line2, = plt.plot([2.8, 3.7], [4.65, 4.65], '--', linewidth=2, color=[0.0, 0.0, 0.0], alpha=0.5)
+    plt.text(4.1, 5.65, 'Walking MDC - Stroke [27]', fontdict=FONT_DICT_SMALL, color=[0.0, 0.0, 0.0], va='center')
+    plt.annotate('', xytext=(4.05, 5.65), xycoords='data', xy=(3.75, 4.65), arrowprops=dict(arrowstyle="->"))
 
     format_axis(plt.gca())
     format_ticks(plt.gca())
-    plt.tight_layout(rect=[0., -0.01, 1, 1.01])
-    plt.legend(list(bars) + [line0, line1, line2], [
-        'GaitDynamics', 'GroundLink [XX]', 'SugaiNet [XX]', 'Running MDC - Healthy [XX]', 'Walking MDC - Healthy [XX]', 'Walking MDC - Stroke [XX]'],
-               frameon=False, fontsize=FONT_SIZE_SMALL, bbox_to_anchor=(0.4, 1.05))       # fontsize=font_size,
+    plt.tight_layout(rect=[-0.03, 0., 1.03, 0.88])
+    plt.legend(list(bars), ['GaitDynamics', 'Convolutional Neural Network [19]', 'Recurrent Neural Network [20]'],
+               frameon=False, fontsize=FONT_SIZE_SMALL, bbox_to_anchor=(0.7, 1.2), ncols=1)
     plt.savefig(f'exports/da_grf.png', dpi=300, bbox_inches='tight')
     plt.show()
 
 
 def draw_fig_3(fast_run=False):
     def format_ticks(ax_plt):
-        ax_plt.set_ylabel('Mean Absolute Error of Peak $f_v$ (\% Body Weight)', fontdict=FONT_DICT_SMALL)
+        ax_plt.text(-0.7, 20, 'Mean Absolute Error of Vertical Force Estimation (\% BW)', rotation=90, fontdict=FONT_DICT_SMALL, verticalalignment='center')
+        ax_plt.text(-1., 37, 'Better', rotation=90, fontdict=FONT_DICT_SMALL, color='green', verticalalignment='center')
+        ax_plt.annotate('', xy=(-0.11, 0.8), xycoords='axes fraction', xytext=(-0.11, 1.),
+                        arrowprops=dict(arrowstyle="->", color='green'))
         ax_plt.set_yticks([0, 10, 20, 30, 40])
         ax_plt.set_yticklabels([0, 10, 20, 30, 40], fontdict=FONT_DICT_SMALL)
         ax_plt.set_ylim([0, 40])
@@ -336,17 +353,17 @@ def format_ticks(ax_plt):
             masked_segments = test_name.split('_')
             for i_segment, segment in enumerate(segment_list):
                 if segment in masked_segments or segment[:-1] in masked_segments:
-                    ax_text.text(i_test+0.14, 7.9 - i_segment*1.1, segment, fontdict=FONT_DICT_SMALL, color=[0.8, 0.8, 0.8], ha='center')
+                    ax_text.text(i_test*0.96+0.35, 7.9 - i_segment*1.1, segment, fontdict=FONT_DICT_SMALL, color=[0.8, 0.8, 0.8], ha='center')
                 else:
-                    ax_text.text(i_test+0.14, 7.9 - i_segment*1.1, segment, fontdict=FONT_DICT_SMALL, ha='center')
+                    ax_text.text(i_test*0.96+0.35, 7.9 - i_segment*1.1, segment, fontdict=FONT_DICT_SMALL, ha='center')
 
     colors = [np.array(x) / 255 for x in [[70, 130, 180], [207, 154, 130]]]        #  [207, 154, 130], [100, 155, 227]
     folder = 'fast' if fast_run else 'full'
     param_of_interest = 'calcn_l_force_vy_max'
-    fig = plt.figure(figsize=(5.5, 4.2))
+    fig = plt.figure(figsize=(7.7, 4.8))
     rc('text', usetex=True)
     plt.rc('font', family='Helvetica')
-    ax_plt = fig.add_axes([0.1, 0.25, 0.87, 0.62])
+    ax_plt = fig.add_axes([0.14, 0.25, 0.83, 0.66])
 
     full_input = get_all_the_metrics(model_key=f'/{folder}/tf_none_diffusion_filling')[param_of_interest]
     line_1, = plt.plot([-0.3, 7.6], [np.mean(full_input), np.mean(full_input)], color=np.array([70, 130, 180])/255, linewidth=LINE_WIDTH, linestyle='--')
@@ -366,8 +383,8 @@ def format_ticks(ax_plt):
     format_axis(plt.gca())
     format_ticks(ax_plt)
     ax_plt.legend(list(bars) + [line_1], [
-        'Partial-Body Kinematics with Inpainting Filling', 'Partial-Body Kinematics with Median Filling', 'Full-Body Kinematics'],
-                  frameon=False, fontsize=FONT_SIZE_SMALL, bbox_to_anchor=(0.05, 0.88), loc='lower left')
+        'Partial-Body Kinematics with Inpainting Filling (GaitDynamics)', 'Partial-Body Kinematics with Median Filling', 'Full-Body Kinematics (GaitDynamics)'],
+                  frameon=False, fontsize=FONT_SIZE_SMALL, bbox_to_anchor=(0., 0.88), loc='lower left')
     plt.savefig(f'exports/da_segment_filling.png', dpi=300, bbox_inches='tight')
     plt.show()
 

figures/da_guided_ts_uhlrich_0.py

@@ -138,8 +138,12 @@ def loop_all(opt):
                 sub_ts_true_original_rate[test_name].append(true_original_rate)
                 sub_ts_pred_original_rate[test_name].append(pred_original_rate)
 
-            height_m_all[test_name] = trial_of_this_win.height_m
-            weight_kg_all[test_name] = trial_of_this_win.weight_kg
+            if test_name not in height_m_all.keys():
+                height_m_all[test_name] = []
+                weight_kg_all[test_name] = []
+            if ('walking' in trial_of_this_win.sub_and_trial_name) and ('ts' not in trial_of_this_win.sub_and_trial_name.lower()):
+                height_m_all[test_name].append(trial_of_this_win.height_m)
+                weight_kg_all[test_name].append(trial_of_this_win.weight_kg)
 
             # if x_times_lumbar_bending > 1:
             #     name_states_dict = {'true': true_val, 'pred': state_pred.detach().numpy()}
@@ -156,8 +160,8 @@ def loop_all(opt):
 
 if __name__ == "__main__":
     opt = parse_opt()
-    opt.n_guided_steps = 3
-    opt.guidance_lr = 0.01
+    opt.n_guided_steps = 5
+    opt.guidance_lr = 0.02
     opt.guide_x_start_the_beginning_step = 1000
     opt.guide_x_start_the_end_step = 0
     opt.checkpoint = os.path.dirname(os.path.realpath(__file__)) + f"/../trained_models/train-{'2560_diffusion'}.pt"