update for srcvol3

app.py

@@ -34,9 +34,9 @@ IMG_SIZE = 256
 CACHE_API_CALLS = False
 os.makedirs(os.path.join(os.getcwd(), 'clips'), exist_ok=True)
 
-onnx_file = hf_hub_download(repo_id="lumos-motion/nextjump", filename="nextjump_256.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+current_model = 'nextjump_speed'
+onnx_file = hf_hub_download(repo_id="lumos-motion/nextjump", filename=f"{current_model}.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
 
-#onnx_file = 'nextjump.onnx'
 
 if torch.cuda.is_available():
     print("Using CUDA")
@@ -236,16 +236,32 @@ def detect_relay_beeps(video_path, event_start, relay_length=30, n_jumpers=4, be
     return starts, ends
 
 
-def inference(in_video, stream_url, start_time, end_time, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay,
+def inference(in_video, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay,
               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()):
-    progress(0, desc="Downloading clip...")
-    if in_video is None:
-        in_video = download_clips(stream_url, os.path.join(os.getcwd(), 'clips'), start_time, end_time)
-    else:  # local uploaded video (still resize with ffmpeg)
-        in_video = download_clips(in_video, os.path.join(os.getcwd(), 'clips'), start_time, end_time)
+    global current_model
+    if model_choice != current_model:
+        current_model = model_choice
+        onnx_file = hf_hub_download(repo_id="lumos-motion/nextjump", filename=f"{current_model}.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+
+
+        if torch.cuda.is_available():
+            print("Using CUDA")
+            providers = [("CUDAExecutionProvider", {"device_id": torch.cuda.current_device(),
+                                                    "user_compute_stream": str(torch.cuda.current_stream().cuda_stream)})]
+            sess_options = ort.SessionOptions()
+            #sess_options.log_severity_level = 0
+            ort_sess = ort.InferenceSession(onnx_file, sess_options=sess_options, providers=providers)
+        else:
+            print("Using CPU")
+            ort_sess = ort.InferenceSession(onnx_file)
+
+        # warmup inference
+        ort_sess.run(None, {'video': np.zeros((4, 64, 3, IMG_SIZE, IMG_SIZE), dtype=np.float32)})
+
+    in_video = download_clips(in_video, os.path.join(os.getcwd(), 'clips'), '00:00:00', '', use_60fps=use_60fps)
     progress(0, desc="Running inference...")
     has_access = False
     if api_call:
@@ -350,14 +366,14 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
             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()
+                periodLength = y1
                 periodicity = y2.squeeze()
                 marks = y3.squeeze()
                 event_type = y4.squeeze()
                 foot_type = y5.squeeze()
                 phase = y6.squeeze()
                 period_lengths[idx:idx+seq_len] += periodLength[:, 0]
-                period_lengths_rope[idx:idx+seq_len] += periodLength[:, 1]
+                #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
@@ -368,7 +384,7 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
             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]
+    #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]
@@ -442,14 +458,14 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
 
     if LOCAL:
         if both_feet:
-            count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
+            count_msg = f"## Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
         else:
-            count_msg = f"## Reps Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
+            count_msg = f"## Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
     else:
         if both_feet:
-            count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Confidence: {total_confidence:.2f}"
+            count_msg = f"## Count (both feet): {count_pred:.1f}, Confidence: {total_confidence:.2f}"
         else:
-            count_msg = f"## Reps Count (one foot): {count_pred:.1f}, Confidence: {total_confidence:.2f}"
+            count_msg = f"## Count (one foot): {count_pred:.1f}, Confidence: {total_confidence:.2f}"
 
     if api_call:
         if CACHE_API_CALLS:
@@ -458,7 +474,8 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                 periodicity_str = np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', '')
                 periodLength_str = np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', '')
                 full_marks_str = np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', '')
-                f.write(f"{stream_url}\t{start_time}\t{end_time}\t{beep_detection_on}\t{event_length}\t{periodicity_str}\t{periodLength_str}\t{full_marks_str}\t{count_pred}\t{total_confidence}\n")
+                f.write(f"{beep_detection_on}\t{event_length}\t{periodicity_str}\t{periodLength_str}\t{full_marks_str}\t{count_pred}\t{total_confidence}\n")
+        
         if count_only_api:
             return f"{count_pred:.2f} (conf: {total_confidence:.2f})"
         else:
@@ -552,10 +569,10 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                     y='jumps per second',
                     #symbol='frame_type',
                     #symbol_map={'frame': 'circle', 'miss': 'circle-open', 'jump': 'triangle-down'},
-                    color='event_type',
+                    color='periodicity',
                     size='jumps_size',
                     size_max=8,
-                    color_continuous_scale=[(t, c) for t, c in zip(event_type_tick_vals, event_type_colors)],
+                    color_continuous_scale='rainbow',
                     range_color=(0,1),
                     title="Jumping speed (jumps-per-second)",
                     trendline='rolling',
@@ -593,11 +610,11 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     )
     # remove white outline from marks
     fig.update_traces(marker_line_width = 0)
-    fig.update_layout(coloraxis_colorbar=dict(
-        tickvals=event_type_tick_vals,
-        ticktext=['single<br>rope', 'double<br>dutch', 'double<br>unders', 'single<br>bounces', 'double<br>bounces', 'triple<br>unders', 'other'],
-        title='event type'
-    ))
+    # fig.update_layout(coloraxis_colorbar=dict(
+    #     tickvals=event_type_tick_vals,
+    #     ticktext=['single<br>rope', 'double<br>dutch', 'double<br>unders', 'single<br>bounces', 'double<br>bounces', 'triple<br>unders', 'other'],
+    #     title='event type'
+    # ))
 
 
     # -pi/2 phase offset to make the bottom of the plot the start of the jump
@@ -674,20 +691,53 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                         histnorm='percent',
                         title="Distribution of jumping speed (jumps-per-second)")
     
-    # make a bar plot of the event type distribution
-
-    bar = px.bar(x=['single rope', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'], 
-                 y=event_type_probs,
-                 template="plotly_dark",
-                 title="Event Type Distribution",
-                 labels={'x': 'event type', 'y': 'probability'},
-                 range_y=[0, 1])
+    # plot the full count and predict a count for 30s, 60s, and 180s if the video is shorter than that
+    count = np.array(count)
+    regression_plot = px.scatter(x=np.arange(len(count)), y=count,
+                                color=periodicity,
+                                color_continuous_scale='rainbow',
+                                title="Count Prediction (Perfect Run)",
+                                template="plotly_dark")
+    regression_plot.update_coloraxes(colorbar=dict(
+        title="Periodicity"))
+    regression_plot.update_traces(marker=dict(size=6, opacity=0.5))
+    regression_plot.update_layout(
+        xaxis_title="Frame",
+        yaxis_title="Count",
+        xaxis=dict(range=[0, len(count)]),
+        yaxis=dict(range=[0, max(count) * 1.2]),
+        showlegend=False,
+    )
+
+    # add 30s, 60s, and 180s predictions
+    pred_count_30s = int(np.median(jumps_per_second[~misses]) * 30)
+    pred_count_60s = int(np.median(jumps_per_second[~misses]) * 60)
+    pred_count_180s = int(np.median(jumps_per_second[~misses]) * 180)
+    # add text to the plot
+    regression_plot.add_annotation(
+        x=0.5,
+        y=0.95,
+        xref="paper",
+        yref="paper",
+        text=f"No-Miss Count (30s): {pred_count_30s}<br>No-Miss Count (60s): {pred_count_60s}<br>No-Miss Count (180s): {pred_count_180s}",
+        showarrow=False,
+        font=dict(
+            size=16,
+            color="white"
+        ),
+        align="center",
+        bgcolor="rgba(0, 0, 0, 0.5)",
+        bordercolor="white",
+        borderwidth=2,
+        borderpad=4,
+        opacity=0.8
+    )
     try:
         os.remove('temp.wav')
     except FileNotFoundError:
         pass
     
-    return in_video, count_msg, fig, fig_phase_spiral, fig_phase_spiral_marks, hist, bar
+    return in_video, count_msg, fig, fig_phase_spiral, fig_phase_spiral_marks, hist, regression_plot
         
 #css = '#phase-spiral {transform: rotate(0.25turn);}\n#phase-spiral-marks {transform: rotate(0.25turn);}'
 with gr.Blocks() as demo:
@@ -697,12 +747,10 @@ with gr.Blocks() as demo:
                                 max_length=300)
     with gr.Row():
         with gr.Column():
-            in_stream_url = gr.Textbox(label="Stream URL", elem_id='stream-url', visible=True)
-            
-            in_stream_start = gr.Textbox(label="Start Time", elem_id='stream-start', visible=True, value='00:00:00')
-            
-            in_stream_start = gr.Textbox(label="Start Time", elem_id='stream-start', visible=True, value='00:00:00')
-            in_stream_end = gr.Textbox(label="End Time", elem_id='stream-end', visible=True)
+            use_60fps = gr.Checkbox(label="Use 60 FPS", elem_id='use-60fps', visible=True)
+            model_choice =  gr.Dropdown(
+                ["nextjump_speed", "nextjump_all"], label="Model Choice", info="For now just speed-only or general model",
+            )
         with gr.Column():
             beep_detection_on = gr.Checkbox(label="Detect Beeps", elem_id='detect-beeps', visible=True)
             event_length = gr.Textbox(label="Expected Event Length (s)", elem_id='event-length', visible=True)
@@ -740,19 +788,19 @@ 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], 
+    run_button.click(demo_inference, [in_video, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
                      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], 
+    api_dummy_button.click(api_inference, [in_video, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay, count_only, api_token], 
                            outputs=[period_length], api_name='inference')
     examples = [
         #['https://hiemdall-dev2.azurewebsites.net/api/clip/clp_vrpWTyjM/mp4', '00:00:00', '00:01:10', True, 60],
-        ['files/wc2023.mp4', '', '00:00:00', '', True, 30, False, '30', '0.2'],
+        ['files/wc2023.mp4', True, 'nextjump_speed', True, 30, False, '30', '0.2'],
         #['https://hiemdall-dev2.azurewebsites.net/api/playlist/rec_rd2FAyUo/vod', '01:24:22', '01:25:35', True, 60]
         #['https://hiemdall-dev2.azurewebsites.net/api/playlist/rec_PY5Ukaua/vod, '00:52:53', '00:55:00', True, 120]
     ]
     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], 
+                inputs=[in_video, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
                 outputs=[out_video, out_text, out_plot, out_phase_spiral, out_phase, out_hist, out_event_type_dist],
                 fn=demo_inference, cache_examples=False)
 

hls_download.py

@@ -1,7 +1,7 @@
 import subprocess
 import os
 
-def download_clips(stream_url, out_dir, start_time, end_time, resize=True):
+def download_clips(stream_url, out_dir, start_time, end_time, resize=True, use_60fps=False):
     # remove all .mp4 files in out_dir to avoid confusion
     if len(os.listdir(out_dir)) > 5:
         for f in os.listdir(out_dir):
@@ -14,7 +14,7 @@ def download_clips(stream_url, out_dir, start_time, end_time, resize=True):
                 '-i', stream_url,
                 '-c:v', 'libx264',
                 '-crf', '23',
-                '-r', '30',
+                '-r', '30' if not use_60fps else '60',
                 '-maxrate', '2M',
                 '-bufsize', '4M',
                 '-vf', f"scale=-2:300",