updated instruction and buttons and progression bar

app.py

@@ -538,8 +538,8 @@ class GaussianDiffusion(nn.Module):
 
 
 class FillingBase:
-    def fill_param(self, windows, diffusion_model_for_filling):
-        return self.filling(windows, diffusion_model_for_filling, self.update_kinematics_and_masks_for_masking_column)
+    def fill_param(self, windows, diffusion_model_for_filling, progress_callback=None):
+        return self.filling(windows, diffusion_model_for_filling, self.update_kinematics_and_masks_for_masking_column, progress_callback=progress_callback)
 
     @staticmethod
     def update_kinematics_and_masks_for_masking_column(windows, samples, i_win, masks):
@@ -565,9 +565,10 @@ class FillingBase:
 
 class DiffusionFilling(FillingBase):
     @staticmethod
-    def filling(windows, diffusion_model_for_filling, windows_update_func):
+    def filling(windows, diffusion_model_for_filling, windows_update_func, progress_callback=None):
         windows = copy.deepcopy(windows)
-        for i_win in range(0, len(windows), opt.batch_size_inference):
+        total_iterations = len(range(0, len(windows), opt.batch_size_inference))
+        for iteration_idx, i_win in enumerate(range(0, len(windows), opt.batch_size_inference)):
             state_true = torch.stack([win.pose for win in windows[i_win:i_win+opt.batch_size_inference]])
             masks = torch.stack([win.mask for win in windows[i_win:i_win+opt.batch_size_inference]])
             cond = torch.ones([6])
@@ -579,6 +580,11 @@ class DiffusionFilling(FillingBase):
             # samples[:, :, opt.kinetic_diffusion_col_loc] = state_true[:, :, opt.kinetic_diffusion_col_loc]
 
             windows = windows_update_func(windows, samples, i_win, masks)
+            
+            # Report progress: 10% for first iteration, 90% for last iteration
+            if progress_callback:
+                progress = 0.1 + 0.8 * (iteration_idx + 1) / total_iterations
+                progress_callback(progress)
         return windows
 
     def __str__(self):
@@ -1238,16 +1244,16 @@ def convertDfToGRFMot(df, out_path, dt, time_column):
 
     plate_num = 2
 
-    for i in range(1, 1 + plate_num):
-        out_file.write('\t' + f'force{i}_vx')
-        out_file.write('\t' + f'force{i}_vy')
-        out_file.write('\t' + f'force{i}_vz')
-        out_file.write('\t' + f'force{i}_px')
-        out_file.write('\t' + f'force{i}_py')
-        out_file.write('\t' + f'force{i}_pz')
-        out_file.write('\t' + f'torque{i}_x')
-        out_file.write('\t' + f'torque{i}_y')
-        out_file.write('\t' + f'torque{i}_z')
+    for i, side in zip(range(1, 1 + plate_num), ['r', 'l']):
+        out_file.write('\t' + f'force_{side}_vx')
+        out_file.write('\t' + f'force_{side}_vy')
+        out_file.write('\t' + f'force_{side}_vz')
+        out_file.write('\t' + f'force_{side}_px')
+        out_file.write('\t' + f'force_{side}_py')
+        out_file.write('\t' + f'force_{side}_pz')
+        out_file.write('\t' + f'torque_{side}_x')
+        out_file.write('\t' + f'torque_{side}_y')
+        out_file.write('\t' + f'torque_{side}_z')
 
     out_file.write('\n')
     for i in range(numFrames):
@@ -2626,78 +2632,6 @@ class TrialData:
 
 """ ============================ End dataset.py ============================ """
 
-
-# def usr_inputs():
-#     opt = parse_opt()
-
-#     # # [DEBUG]
-#     # opt.height_m = 1.84
-#     # opt.weight_kg = 71.4
-#     # opt.treadmill_speed = 1.15
-#     # opt.subject_osim_model = opt.subject_data_path + '/Scaled_generic_no_arm.osim'
-
-#     input("Upload .mot files and a .osim file to Colab via \n \
-#     1) click the folder button on the left side of the window \n \
-#     2) upload .mot files to the default folder (/content/) by clicking \"Upload to session storage\" button and \n \
-#     3) upload .osim file to the default folder (/content/) by clicking \"Upload to session storage\" button. \n \
-#     Please confirm completion and then entering anything ")
-
-#     file_paths = []
-#     for file in os.listdir(opt.subject_data_path):
-#         file_path = os.path.join(opt.subject_data_path, file)
-#         if file.endswith(".mot") and '_pred___' not in file:
-#             file_paths.append(file_path)
-#     if len(file_paths) == 0:
-#         raise RuntimeError(f'No .mot file found. Upload the .mot file to Colab')
-#     opt.file_paths = file_paths
-
-#     osim_paths = []
-#     for file in os.listdir(opt.subject_data_path):
-#         file_path = os.path.join(opt.subject_data_path, file)
-#         if file.endswith(".osim"):
-#             osim_paths.append(file_path)
-#     if len(osim_paths) > 1:
-#         print(f'Multiple .osim files found.')
-#         [print(f'{i}: {file_path}') for i, file_path in enumerate(osim_paths)]
-#         i_file = int(input(f'Choose the .osim file entering its index (between 0 and {len(osim_paths)-1}): '))
-#         opt.subject_osim_model = osim_paths[i_file]
-#     elif len(osim_paths) == 0:
-#         raise RuntimeError(f'No .osim file found. Upload the .osim file to Colab')
-#     else:
-#         opt.subject_osim_model = osim_paths[0]
-#     print()
-
-#     while True:
-#         try:
-#             opt.height_m = float(input("Enter the subject's height in meter: "))
-#             if opt.height_m > 2.5 or opt.height_m < 1:
-#                 raise ValueError()
-#             print()
-#             break
-#         except ValueError:
-#             print("Invalid input. Please enter a floating-point number between 1.0 and 2.5.")
-#     while True:
-#         try:
-#             opt.weight_kg = float(input("Enter the subject's weight in Kg: "))
-#             if opt.weight_kg > 200 or opt.weight_kg < 30:
-#                 raise ValueError()
-#             print()
-#             break
-#         except ValueError:
-#             print("Invalid input. Please enter a floating-point number between 30 and 200.")
-#     while True:
-#         try:
-#             opt.treadmill_speed = float(input("Enter treadmill speed in m/s (enter 0 for overground gait): "))
-#             if opt.treadmill_speed > 10 or opt.treadmill_speed < 0:
-#                 raise ValueError()
-#             print()
-#             break
-#         except ValueError:
-#             print("Invalid input. Please enter a floating-point number between 0 and 10.")
-
-#     return opt
-
-
 def predict_grf_and_missing_kinematics():
     refinement_model = BaselineModel(opt, TransformerEncoderArchitecture)
     dataset = MotionDataset(opt, normalizer=refinement_model.normalizer)
@@ -2756,7 +2690,7 @@ import io
 import shutil
 
 opt = parse_opt()
-def gradio_predict(mot_file, osim_file, height_m, weight_kg, treadmill_speed):
+def gradio_predict(mot_file, osim_file, height_m, weight_kg, treadmill_speed, progress_callback=None):
     """Gradio interface function for GRF prediction."""
     if mot_file is None:
         return "Please upload a .mot file"
@@ -2790,7 +2724,7 @@ def gradio_predict(mot_file, osim_file, height_m, weight_kg, treadmill_speed):
         opt.subject_osim_model = osim_path
         
         # Run prediction with updated opt
-        output_files = predict_grf_and_missing_kinematics_with_opt()
+        output_files = predict_grf_and_missing_kinematics_with_opt(progress_callback=progress_callback)
         
         if output_files:
             return "Prediction completed! Download files below.", output_files[0], output_files[1] if len(output_files) > 1 else None
@@ -2799,7 +2733,7 @@ def gradio_predict(mot_file, osim_file, height_m, weight_kg, treadmill_speed):
     except Exception as e:
         return f"Error during prediction: {str(e)}", None, None
 
-def predict_grf_and_missing_kinematics_with_opt():
+def predict_grf_and_missing_kinematics_with_opt(progress_callback=None):
     """Modified prediction function that accepts opt as parameter."""
     refinement_model = BaselineModel(opt, TransformerEncoderArchitecture)
     dataset = MotionDataset(opt, normalizer=refinement_model.normalizer)
@@ -2816,7 +2750,7 @@ def predict_grf_and_missing_kinematics_with_opt():
                   f'\nGenerating missing kinematics for {dataset.file_names[i_trial]}')
             if diffusion_model_for_filling is None:
                 diffusion_model_for_filling, _ = load_diffusion_model(opt)
-            windows_reconstructed = filling_method.fill_param(windows, diffusion_model_for_filling)
+            windows_reconstructed = filling_method.fill_param(windows, diffusion_model_for_filling, progress_callback=progress_callback)
         else:
             windows_reconstructed = windows
 
@@ -2932,43 +2866,74 @@ with gr.Blocks(css="""
     }
     #plotting_panel {min-height: 700px !important;}
 """) as demo:
-    gr.Markdown("""
-                # GaitDynamics - Ground Reaction Force and Kinematics Prediction
-                This code is for ground reaction force and missing kinematics prediction using flexible combinations of OpenSim joint angles. The joint angles should meet the following criteria:<br>
-                &nbsp;&nbsp;&nbsp;&nbsp;1&#46; Use OpenSim <a href="https://www.google.com/url?q=https%3A%2F%2Fdrive.google.com%2Fdrive%2Ffolders%2F12CStQmsNVx_UEr9mkBkUApKZgjKOfGUq">Rajagopal Model without Arms</a>.<br>
-                &nbsp;&nbsp;&nbsp;&nbsp;2&#46; Resample to 100 Hz if you have a different sampling rate.<br>
-                To use this code:<br>
-                &nbsp;&nbsp;&nbsp;&nbsp;1&#46; Upload .osim and .mot files. Here are example 
-                <a href="https://github.com/stanfordnmbl/GaitDynamics/blob/f7723a0649d3983f9fd939e89de6fddcb5ea3504/example_usage/example_opensim_model.osim">.osim</a>, 
-                <a href="https://github.com/stanfordnmbl/GaitDynamics/blob/f7723a0649d3983f9fd939e89de6fddcb5ea3504/example_usage/example_mot_complete_kinematics.mot">.mot complete kinematics</a>, 
-                and <a href="https://github.com/stanfordnmbl/GaitDynamics/blob/f7723a0649d3983f9fd939e89de6fddcb5ea3504/example_usage/example_mot_missing_knee_kinematics.mot">.mot incomplete kinematics</a>.<br>
-                &nbsp;&nbsp;&nbsp;&nbsp;2&#46; Enter the height, weight, and treadmill speed, respectively.<br>
-                &nbsp;&nbsp;&nbsp;&nbsp;3&#46; Click " Submit " button."""
-                )
+    gr.Markdown(
+        """
+        # GaitDynamics - Ground Reaction Force and Kinematics Prediction
+        ### General instructions
+        This code is for ground reaction force and missing kinematics prediction using flexible combinations of OpenSim joint angles. The joint angles should meet the following criteria:
+        <div style="padding-left: 2em; text-indent: -1.0em;">1. Use OpenSim <a href="https://simtk.org/projects/full_body">Rajagopal Model without Arms</a>.</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">2. Resample to 100 Hz if you have a different sampling rate.</div>
+        
+        To use this code:
+        <div style="padding-left: 2em; text-indent: -1.0em;">1. Upload OpenSim model (.osim) and kinematics (.mot) file. If kinematics are provided for all coordinates, the model will only predict ground reaction forces. If kinematics are missing, the model will also predict missing kinematics.</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">2. Enter the height and weight of the participant, and the treadmill speed (if applicable).</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">3. Click the "Submit " button.</div>
+        
+        ### Processing example data
+        <div style="padding-left: 2em; text-indent: -1.0em;">1. Download the example files:</div>
+        """
+    )
+    with gr.Column():
+        with gr.Row():
+            gr.DownloadButton("a. OpenSim model file (.osim)", value="example_opensim_model.osim", size="sm", scale=1)
+            gr.DownloadButton("b. complete kinematics file with all coordinates (.mot)", value="example_mot_complete_kinematics.mot", size="sm", scale=1)
+            gr.DownloadButton("c. incomplete kinematics file with missing knee coordinate kinematics (.mot)", value="example_mot_missing_knee_kinematics.mot", size="sm", scale=1)
+    gr.Markdown(
+        """
+        <div style="padding-left: 2em; text-indent: -1.0em;">2. Upload data for either of the two example .mot files.</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">a. Predict ground reaction forces with a complete kinematics file by uploading the model file (.osim) and complete kinematics file (.mot).</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">b. Predict ground reaction forces with an incomplete kinematics file by uploading the model file (.osim) and the incomplete kinematics file (.mot).</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">3. Update the following input parameters.</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">a. Height = 1.84 m</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">b. Weight = 70 kg</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">c. Treadmill speed = 1.15 m/s</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">4. Click "Submit".</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">a. The example with complete kinematics should take a few seconds to complete.</div>
+        <div style="padding-left: 3.5em; text-indent: -1.0em;">b. The example with incomplete kinematics should take a few minutes to complete.</div>
+        <div style="padding-left: 2em; text-indent: -1.0em;">5. Plot the data, starting with GRF Results and choosing columns (e.g., force_r_vy).</div>
+        """
+    )
 
     gr.Markdown("### Input Data and Parameters")
     with gr.Row():
         mot_in = gr.File(label="Upload .mot file", file_types=[".mot"])
         osim_in = gr.File(label="Upload .osim file", file_types=[".osim"])
     with gr.Row():
-        height_in = gr.Number(label="Height (meters)", value=1.75, minimum=1.0, maximum=2.5, step=0.01)
-        weight_in = gr.Number(label="Weight (kg)", value=70, minimum=30, maximum=200, step=0.1)
-        speed_in = gr.Number(label="Treadmill speed (m/s, 0 for overground)", value=0, minimum=0, maximum=10, step=0.01)
+        height_in = gr.Number(label="Height (meters)", value=1.83, minimum=1.0, maximum=2.5, step=0.01)
+        weight_in = gr.Number(label="Weight (kg)", value=71.4, minimum=30, maximum=200, step=0.1)
+        speed_in = gr.Number(label="Treadmill speed (m/s, 0 for overground)", value=1.15, minimum=0, maximum=10, step=0.01)
 
     submit_btn = gr.Button("Submit")
 
     gr.Markdown("### Results")
     with gr.Row():
         status_out = gr.Textbox(label="Status")
+    grf_path = gr.State(None)
+    kin_path = gr.State(None)
     with gr.Row():
-        grf_file_out = gr.File(label="Download Ground Reaction Force Results", interactive=False, elem_classes=["download-file"])
-        kin_file_out = gr.File(label="Download Missing Kinematics (if applicable)", interactive=False, elem_classes=["download-file"])
+        grf_file_out = gr.DownloadButton("Download Ground Reaction Force Results", size="sm", scale=1, visible=False)
+        kin_file_out = gr.DownloadButton("Download Missing Kinematics (0 for overground)", size="sm", scale=1, visible=False)
 
     # Plotting panel
     with gr.Column(elem_id="plotting_panel"):
         gr.Markdown("""
                     ### Visualization
-                    Select one of the generated output files and a column to visualize over time.
+                    <div style="padding-left: 2em; text-indent: -1.0em;">Ground reaction force (GRF) results are available for all inputs. Missing kinematics results are only available if incomplete kinematics were used as inputs.</div>
+                    <div style="padding-left: 2em; text-indent: -1.0em;">GRF column names follow these conventions:</div>
+                    <div style="padding-left: 3.5em; text-indent: -1.0em;">• Starts with force_l or force_r: indicates the left or right foot, respectively</div>
+                    <div style="padding-left: 3.5em; text-indent: -1.0em;">• Ends in _vx, _vy, _vz: indicates the magnitude value of the force in the x, y, and z directions, respectively.</div>
+                    <div style="padding-left: 3.5em; text-indent: -1.0em;">• Ends in _vx, _vy, _vz: indicates the point location of the center of pressure in the x, y, and z directions, respectively.</div>
+                    <div style="padding-left: 3.5em; text-indent: -1.0em;">• x is the anterior direction, y is the vertical direction, and z is the direction pointing to the right (medial or lateral direction depending on the leg).</div>
                     """)
         with gr.Row():
             file_select = gr.Radio(choices=["GRF Results", "Missing Kinematics"], value="GRF Results", label="Select output file")
@@ -2977,33 +2942,35 @@ with gr.Blocks(css="""
         plot_out = gr.Plot(label="Time series plot")
 
     # Wire interactions
+    def enhanced_predict(mot_in, osim_in, height_in, weight_in, speed_in, current_select, current_grf, current_kin, progress=gr.Progress()):
+        progress(0.0, desc="Starting...")
+        
+        def progress_callback(pct):
+            progress(pct, desc=f"Processing diffusion model... {int(pct*100)}%")
+        
+        status, new_grf, new_kin = gradio_predict(mot_in, osim_in, height_in, weight_in, speed_in, progress_callback=progress_callback)
+        progress(1.0, desc="Complete!")
+        
+        columns = list_mot_columns(current_select, new_grf, new_kin)
+        grf_visible = new_grf is not None
+        kin_visible = new_kin is not None
+        return status, new_grf, new_kin, gr.update(value=new_grf, visible=grf_visible), gr.update(value=new_kin, visible=kin_visible), columns
+
     submit_btn.click(
-        fn=gradio_predict,
-        inputs=[mot_in, osim_in, height_in, weight_in, speed_in],
-        outputs=[status_out, grf_file_out, kin_file_out]
+        fn=enhanced_predict,
+        inputs=[mot_in, osim_in, height_in, weight_in, speed_in, file_select, grf_path, kin_path],
+        outputs=[status_out, grf_path, kin_path, grf_file_out, kin_file_out, column_select]
     )
 
     file_select.change(
         fn=list_mot_columns,
-        inputs=[file_select, grf_file_out, kin_file_out],
-        outputs=[column_select]
-    )
-
-    grf_file_out.change(
-        fn=list_mot_columns,
-        inputs=[file_select, grf_file_out, kin_file_out],
-        outputs=[column_select]
-    )
-
-    kin_file_out.change(
-        fn=list_mot_columns,
-        inputs=[file_select, grf_file_out, kin_file_out],
+        inputs=[file_select, grf_path, kin_path],
         outputs=[column_select]
     )
 
     plot_btn.click(
         fn=plot_mot_signal,
-        inputs=[file_select, column_select, grf_file_out, kin_file_out],
+        inputs=[file_select, column_select, grf_path, kin_path],
         outputs=[plot_out]
     )