Update app.py

app.py

@@ -30,11 +30,14 @@ from hls_download import download_clips
 #plt.style.use('dark_background')
 
 LOCAL = False
-IMG_SIZE = 256
+IMG_SIZE = 192
 CACHE_API_CALLS = True
 os.makedirs(os.path.join(os.getcwd(), 'clips'), exist_ok=True)
 
-onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="nextjump.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+onnx_file = hf_hub_download(repo_id="lumos-motion/nextjump", filename="nextjump_192.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+
+#onnx_file = 'nextjump.onnx'
+
 if torch.cuda.is_available():
     print("Using CUDA")
     providers = [("CUDAExecutionProvider", {"device_id": torch.cuda.current_device(),
@@ -234,8 +237,7 @@ def detect_relay_beeps(video_path, event_start, relay_length=30, n_jumpers=4, be
 
 
 def inference(in_video, stream_url, start_time, end_time, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay,
-              count_only_api, api_key,
-              img_size=256, seq_len=64, stride_length=32, stride_pad=3, batch_size=4,
+              count_only_api, api_key, seq_len=64, stride_length=32, stride_pad=3, batch_size=4,
               miss_threshold=0.8, marks_threshold=0.5, median_pred_filter=True, both_feet=True, 
               api_call=False,
               progress=gr.Progress()):
@@ -269,7 +271,7 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     seconds = length / fps
     all_frames = []
     frame_i = 0
-    resize_amount = max((img_size + 64) / frame_width, (img_size + 64) / frame_height)
+    resize_amount = max((IMG_SIZE + 64) / frame_width, (IMG_SIZE + 64) / frame_height)
     while cap.isOpened():
         frame_i += 1
         
@@ -286,18 +288,21 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         frame = cv2.resize(frame, (0, 0), fx=resize_amount, fy=resize_amount, interpolation=cv2.INTER_CUBIC)
         frame_center_x = frame.shape[1] // 2
         frame_center_y = frame.shape[0] // 2
-        crop_x = frame_center_x - img_size // 2
-        crop_y = frame_center_y - img_size // 2
-        frame = frame[crop_y:crop_y+img_size, crop_x:crop_x+img_size]
+        crop_x = frame_center_x - IMG_SIZE // 2
+        crop_y = frame_center_y - IMG_SIZE // 2
+        frame = frame[crop_y:crop_y+IMG_SIZE, crop_x:crop_x+IMG_SIZE]
         all_frames.append(frame)
         
     cap.release()
 
     length = len(all_frames)
     period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
+    period_lengths_rope = np.zeros(len(all_frames) + seq_len + stride_length)
     periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
     full_marks = np.zeros(len(all_frames) + seq_len + stride_length)
     event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 7))
+    phase_sin = np.zeros(len(all_frames) + seq_len + stride_length)
+    phase_cos = np.zeros(len(all_frames) + seq_len + stride_length)
     period_length_overlaps = np.zeros(len(all_frames) + seq_len + stride_length)
     event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 7))
     for _ in range(seq_len + stride_length):  # pad full sequence
@@ -309,7 +314,7 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         for i in range(0, length + stride_length - stride_pad, stride_length):
             batch = all_frames[i:i + seq_len]
             Xlist = []
-            preprocess_tasks = [(idx, executor.submit(preprocess_image, img, img_size)) for idx, img in enumerate(batch)]
+            preprocess_tasks = [(idx, executor.submit(preprocess_image, img, IMG_SIZE)) for idx, img in enumerate(batch)]
             for idx, future in sorted(preprocess_tasks, key=lambda x: x[0]):
                 Xlist.append(future.result())
             
@@ -342,23 +347,35 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
             y2_out = outputs[1]
             y3_out = outputs[2]
             y4_out = outputs[3]
-            for y1, y2, y3, y4, idx in zip(y1_out, y2_out, y3_out, y4_out, idx_list):
+            y5_out = outputs[4]
+            y6_out = outputs[5]
+            for y1, y2, y3, y4, y5, y6, idx in zip(y1_out, y2_out, y3_out, y4_out, y5_out, y6_out, idx_list):
                 periodLength = y1.squeeze()
                 periodicity = y2.squeeze()
                 marks = y3.squeeze()
                 event_type = y4.squeeze()
-                period_lengths[idx:idx+seq_len] += periodLength
+                foot_type = y5.squeeze()
+                phase = y6.squeeze()
+                period_lengths[idx:idx+seq_len] += periodLength[:, 0]
+                period_lengths_rope[idx:idx+seq_len] += periodLength[:, 1]
                 periodicities[idx:idx+seq_len] += periodicity
                 full_marks[idx:idx+seq_len] += marks
                 event_type_logits[idx:idx+seq_len] += event_type
+                phase_sin[idx:idx+seq_len] += phase[:, 1]
+                phase_cos[idx:idx+seq_len] += phase[:, 0]
                 period_length_overlaps[idx:idx+seq_len] += 1
                 event_type_logit_overlaps[idx:idx+seq_len] += 1
             del y1_out, y2_out, y3_out, y4_out  # free up memory
             
     periodLength = np.divide(period_lengths, period_length_overlaps, where=period_length_overlaps!=0)[:length]
+    periodLength_rope = np.divide(period_lengths_rope, period_length_overlaps, where=period_length_overlaps!=0)[:length]
     periodicity = np.divide(periodicities, period_length_overlaps, where=period_length_overlaps!=0)[:length]
     full_marks = np.divide(full_marks, period_length_overlaps, where=period_length_overlaps!=0)[:length]
     per_frame_event_type_logits = np.divide(event_type_logits, event_type_logit_overlaps, where=event_type_logit_overlaps!=0)[:length]
+    phase_sin = np.divide(phase_sin, period_length_overlaps, where=period_length_overlaps!=0)[:length]
+    # negate sin to make the bottom of the plot the start of the jump
+    phase_sin = -phase_sin
+    phase_cos = np.divide(phase_cos, period_length_overlaps, where=period_length_overlaps!=0)[:length]
     event_type_logits = np.mean(per_frame_event_type_logits, axis=0)
     # softmax of event type logits  
     event_type_probs = np.exp(event_type_logits) / np.sum(np.exp(event_type_logits))
@@ -469,33 +486,44 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
 
    
 
-    fig, axs = plt.subplots(4, 1, figsize=(12, 10)) # Added a plot for count
+    # fig, axs = plt.subplots(5, 1, figsize=(14, 10)) # Added a plot for count
 
-    # Ensure data exists before plotting
-    axs[0].plot(periodLength)
-    axs[0].set_title(f"Stream 0 - Period Length")
+    # # Ensure data exists before plotting
+    # axs[0].plot(periodLength, label='Period Length')
+    # axs[0].plot(periodLength_rope, label='Period Length (Rope)')
+    # axs[0].set_title(f"Stream 0 - Period Length")
+    # axs[0].legend()
 
-    axs[1].plot(periodicity)
-    axs[1].set_title("Stream 0 - Periodicity")
-    axs[1].set_ylim(0, 1)
-    axs[1].axhline(miss_threshold, color='r', linestyle=':', label=f'Miss Thresh ({miss_threshold})')
+    # axs[1].plot(periodicity)
+    # axs[1].set_title("Stream 0 - Periodicity")
+    # axs[1].set_ylim(0, 1)
+    # axs[1].axhline(miss_threshold, color='r', linestyle=':', label=f'Miss Thresh ({miss_threshold})')
 
 
-    axs[2].plot(full_marks, label='Raw Marks')
-    marks_peaks_vis, _ = find_peaks(full_marks, distance=3, height=marks_threshold)
-    axs[2].plot(marks_peaks_vis, np.array(full_marks)[marks_peaks_vis], "x", label='Detected Peaks')
-    axs[2].set_title("Stream 0 - Marks")
-    axs[2].set_ylim(0, 1)
-    axs[2].axhline(marks_threshold, color='r', linestyle=':', label=f'Mark Thresh ({marks_threshold})')
+    # axs[2].plot(full_marks, label='Raw Marks')
+    # marks_peaks_vis, _ = find_peaks(full_marks, distance=3, height=marks_threshold)
+    # axs[2].plot(marks_peaks_vis, np.array(full_marks)[marks_peaks_vis], "x", label='Detected Peaks')
+    # axs[2].set_title("Stream 0 - Marks")
+    # axs[2].set_ylim(0, 1)
+    # axs[2].axhline(marks_threshold, color='r', linestyle=':', label=f'Mark Thresh ({marks_threshold})')
 
+    # # plot phase
+    # axs[3].plot(phase_sin, label='Phase Sin')
+    # axs[3].plot(phase_cos, label='Phase Cos')
+    # axs[3].set_title("Stream 0 - Phase")
+    # axs[3].set_ylim(-1, 1)
+    # axs[3].axhline(0, color='r', linestyle=':', label='Zero Line')
+    # axs[3].legend()
 
-    axs[3].plot(count)
-    axs[3].set_title("Stream 0 - Calculated Count")
 
-    plt.tight_layout()
+    # axs[4].plot(count)
+    # axs[4].set_title("Stream 0 - Calculated Count")
 
-    plt.savefig('plot.png')
-    plt.close()
+    # plt.tight_layout()
+
+    # plt.savefig('plot.png')
+    # plt.close()
+    
 
     jumps_per_second = np.clip(1 / ((periodLength / fps) + 0.0001), 0, 10)
     jumping_speed = np.copy(jumps_per_second)
@@ -508,6 +536,8 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                                  'jumping speed': jumping_speed,
                                 'jumps per second': jumps_per_second,
                                 'periodicity': periodicity,
+                                'phase sin': phase_sin,
+                                'phase cos': phase_cos,
                                 'miss': misses,
                                 'frame_type': frame_type,
                                 'event_type': per_frame_event_types,
@@ -569,6 +599,74 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         title='event type'
     ))
 
+
+    # -pi/2 phase offset to make the bottom of the plot the start of the jump
+    # phase_sin = np.sin(np.arctan2(phase_sin, phase_cos) - np.pi / 2)
+    # phase_cos = np.cos(np.arctan2(phase_sin, phase_cos) - np.pi / 2)
+    
+    # plot phase spiral using plotly
+    fig_phase_spiral = px.scatter(x=phase_cos, y=phase_sin,
+                                color=jumps_per_second,
+                                color_continuous_scale='plasma',
+                                title="Phase Spiral (speed)",
+                                template="plotly_dark")
+    fig_phase_spiral.update_traces(marker=dict(size=4, opacity=0.5))
+    fig_phase_spiral.update_layout(
+        xaxis_title="Phase Cos",
+        yaxis_title="Phase Sin",
+        xaxis=dict(range=[-1, 1]),
+        yaxis=dict(range=[-1, 1]),
+        showlegend=False,
+    )
+    # label colorbar as time
+    fig_phase_spiral.update_coloraxes(colorbar=dict(
+        title="Jumps per second"))
+    # make axes equal
+    fig_phase_spiral.update_layout(
+        xaxis=dict(scaleanchor="y"),
+        yaxis=dict(constrain="domain"),
+    )
+    # overlay line plot of phase sin and cos
+    fig_phase_spiral.add_traces(px.line(x=phase_cos, y=phase_sin).data)
+    fig_phase_spiral.update_traces(line=dict(width=0.5, color='rgba(255, 255, 255, 0.25)'))
+
+    # plot phase consistency (sin^2 + cos^2 = 1) as a line plot
+    # phase_consistency = phase_sin**2 + phase_cos**2
+    # #phase_consistency = medfilt(phase_consistency, 5)
+    # fig_phase = px.line(x=np.linspace(0, 1, len(phase_sin)), y=phase_consistency,
+    #                     title="Phase Consistency (sin^2 + cos^2)",
+    #                     labels={'x': 'Frame', 'y': 'Phase Consistency'},
+    #                     template="plotly_dark")
+
+    # plot phase spiral colored by mark_preds
+    fig_phase_spiral_marks = px.scatter(x=phase_cos, y=phase_sin,
+                                color=full_marks,
+                                color_continuous_scale='Jet',
+                                title="Phase Spiral (marks)",
+                                template="plotly_dark")
+    fig_phase_spiral_marks.update_traces(marker=dict(size=4, opacity=0.5))
+    fig_phase_spiral_marks.update_layout(
+        xaxis_title="Phase Cos",
+        yaxis_title="Phase Sin",
+        xaxis=dict(range=[-1, 1]),
+        yaxis=dict(range=[-1, 1]),
+        showlegend=False,
+    )
+    # label colorbar as time
+    fig_phase_spiral_marks.update_coloraxes(colorbar=dict(
+        title="Marks"))
+    # make axes equal
+    fig_phase_spiral_marks.update_layout(
+        xaxis=dict(scaleanchor="y"),
+        yaxis=dict(constrain="domain"),
+    )
+    # overlay line plot of phase sin and cos
+    fig_phase_spiral_marks.add_traces(px.line(x=phase_cos, y=phase_sin).data)
+    fig_phase_spiral_marks.update_traces(line=dict(width=0.5, color='rgba(255, 255, 255, 0.25)'))
+    
+
+    
+
     hist = px.histogram(df, 
                         x="jumps per second", 
                         template="plotly_dark", 
@@ -589,9 +687,9 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     except FileNotFoundError:
         pass
     
-    return in_video, count_msg, fig, hist, bar
+    return in_video, count_msg, fig, fig_phase_spiral, fig_phase_spiral_marks, hist, bar
         
-
+#css = '#phase-spiral {transform: rotate(0.25turn);}\n#phase-spiral-marks {transform: rotate(0.25turn);}'
 with gr.Blocks() as demo:
     with gr.Row():
         in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', 
@@ -628,6 +726,11 @@ with gr.Blocks() as demo:
                 periodicity = gr.Textbox(label="Periodicity", elem_id='periodicity', visible=False)
         with gr.Row():
             out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
+        with gr.Row():
+            with gr.Column():
+                out_phase_spiral = gr.Plot(label="Phase Spiral", elem_id='phase-spiral')
+            with gr.Column():
+                out_phase = gr.Plot(label="Phase Sin/Cos", elem_id='phase-spiral-marks')
         with gr.Row():
             with gr.Column():
                 out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
@@ -638,7 +741,7 @@ with gr.Blocks() as demo:
     demo_inference = partial(inference, count_only_api=False, api_key=None)
     
     run_button.click(demo_inference, [in_video, in_stream_url, in_stream_start, in_stream_end, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
-                     outputs=[out_video, out_text, out_plot, out_hist, out_event_type_dist])
+                     outputs=[out_video, out_text, out_plot, out_phase_spiral, out_phase, out_hist, out_event_type_dist])
     api_inference = partial(inference, api_call=True)
     api_dummy_button.click(api_inference, [in_video, in_stream_url, in_stream_start, in_stream_end, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay, count_only, api_token], 
                            outputs=[period_length], api_name='inference')
@@ -650,7 +753,7 @@ with gr.Blocks() as demo:
     ]
     gr.Examples(examples, 
                 inputs=[in_video, in_stream_url, in_stream_start, in_stream_end, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
-                outputs=[out_video, out_text, out_plot, out_hist, out_event_type_dist],
+                outputs=[out_video, out_text, out_plot, out_phase_spiral, out_phase, out_hist, out_event_type_dist],
                 fn=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces')