Update app.py

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import os
 import cv2
 import math
+import json
 import subprocess
 import matplotlib
 matplotlib.use('Agg')
@@ -26,28 +27,22 @@ from huggingface_hub import HfApi
 
 from hls_download import download_clips
 
-plt.style.use('dark_background')
+#plt.style.use('dark_background')
 
-LOCAL = False
 LOCAL = False
 IMG_SIZE = 256
 CACHE_API_CALLS = True
 os.makedirs(os.path.join(os.getcwd(), 'clips'), exist_ok=True)
-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'])
 if torch.cuda.is_available():
-    print("Using CUDA")
     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
-    #sess_options.log_severity_level = 0
     ort_sess = ort.InferenceSession(onnx_file, sess_options=sess_options, providers=providers)
 else:
-    print("Using CPU")
     print("Using CPU")
     ort_sess = ort.InferenceSession(onnx_file)
 
@@ -245,10 +240,6 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
               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)
     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)
@@ -448,13 +439,58 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         if count_only_api:
             return f"{count_pred:.2f} (conf: {total_confidence:.2f})"
         else:
-            return np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''), \
-                np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''), \
-                np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''), \
-                f"reps: {count_pred:.2f}, marks: {marks_count_pred:.1f}, confidence: {total_confidence:.2f}", \
-                f"single_rope_speed: {event_type_probs[0]:.3f}, double_dutch: {event_type_probs[1]:.3f}, double_unders: {event_type_probs[2]:.3f}, single_bounce: {event_type_probs[3]:.3f}"
+            # create a nice json object to return
+            results_dict = {
+                "periodLength": np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''),
+                "periodicity": np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''),
+                "full_marks": np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''),
+                "cum_count": np.array2string(np.array(count), formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''),
+                "count": f"{count_pred:.2f}",
+                "marks": f"{marks_count_pred:.1f}",
+                "confidence": f"{total_confidence:.2f}",
+                "single_rope_speed": f"{event_type_probs[0]:.3f}",
+                "double_dutch": f"{event_type_probs[1]:.3f}",
+                "double_unders": f"{event_type_probs[2]:.3f}",
+                "single_bounce": f"{event_type_probs[3]:.3f}"
+            }
+            if beep_detection_on:
+                results_dict['event_start'] = event_start
+                results_dict['event_end'] = event_end
+            if relay_detection_on:
+                results_dict['relay_starts'] = relay_starts
+                results_dict['relay_ends'] = relay_ends
+            return json.dumps(results_dict)
+
    
 
+    fig, axs = plt.subplots(4, 1, figsize=(12, 10)) # Added a plot for count
+
+    # Ensure data exists before plotting
+    axs[0].plot(periodLength)
+    axs[0].set_title(f"Stream 0 - Period Length")
+
+    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[3].plot(count)
+    axs[3].set_title("Stream 0 - Calculated Count")
+
+    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)
     misses = periodicity < miss_threshold
@@ -602,7 +638,7 @@ with gr.Blocks() as demo:
                            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],
-        [None, 'https://hiemdall-dev2.azurewebsites.net/api/playlist/rec_UGEhqlMh/vod', '00:00:18', '00:00:55', True, 30],
+        ['files/wc2023.mp4', '', '00:00:00', '', True, 30, False, '', '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]
     ]
@@ -613,14 +649,6 @@ with gr.Blocks() as demo:
 
 
 if __name__ == "__main__":
-    if LOCAL:
-        demo.queue(api_open=True, max_size=15).launch(server_name="0.0.0.0", 
-                                                    server_port=7860, 
-                                                    debug=False,
-                                                    ssl_verify=False,
-                                                    share=False)
-    else:
-        demo.queue(api_open=True, max_size=15).launch(share=False)
     if LOCAL:
         demo.queue(api_open=True, max_size=15).launch(server_name="0.0.0.0", 
                                                     server_port=7860,