Update app.py

app.py

@@ -28,20 +28,26 @@ from hls_download import download_clips
 
 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)
 
@@ -117,12 +123,132 @@ def detect_beeps(video_path, event_length=30, beep_height=0.8):
     return event_start, event_end
 
 
-def inference(in_video, stream_url, start_time, end_time, beep_detection_on, event_length, count_only_api, api_key,
+def detect_relay_beeps(video_path, event_start, relay_length=30, n_jumpers=4, beep_height=0.8):
+    reference_file = 'relay_beep.WAV'
+    fs, beep = wavfile.read(reference_file)
+    beep = beep[:, 0] + beep[:, 1]  # combine stereo to mono
+    video = cv2.VideoCapture(video_path)
+    try:
+        os.remove('temp.wav')
+    except FileNotFoundError:
+        pass
+    audio_convert_command = f'ffmpeg -i {video_path} -vn -acodec pcm_s16le -ar {fs} -ac 2 temp.wav'
+    subprocess.call(audio_convert_command, shell=True)
+    length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = int(video.get(cv2.CAP_PROP_FPS))
+    audio = wavfile.read('temp.wav')[1]
+    audio = (audio[:, 0] + audio[:, 1]) / 2  # combine stereo to mono
+    corr = correlate(audio, beep, mode='same') / audio.size
+    # min max scale to -1, 1
+    corr = 2 * (corr - np.min(corr)) / (np.max(corr) - np.min(corr)) - 1
+    
+    # Calculate total event length in frames
+    total_event_length_frames = fps * relay_length * n_jumpers
+    print(event_start, total_event_length_frames)
+    expected_event_end = event_start + total_event_length_frames
+    
+    # Find all significant peaks in the correlation
+    peaks, _ = find_peaks(corr, height=beep_height, distance=fs)
+    
+    # Convert peaks from sample indices to frame indices
+    peak_frames = [int(peak / fs * fps) for peak in peaks]
+    
+    # For debugging
+    plt.plot(corr)
+    plt.plot(peaks, corr[peaks], "x")
+    plt.savefig('beep.png')
+    plt.close()
+
+    starts = []
+    ends = []
+    
+    # Add the event start for the first jumper
+    starts.append(event_start)
+    
+    # Convert event_start back to sample index for comparison
+    event_start_sample = int(event_start * fs / fps)
+    
+    # Find peaks that come after the event start but before the expected end
+    # Convert expected_event_end to sample index
+    expected_event_end_sample = int(expected_event_end * fs / fps)
+    relevant_peaks = [p for p in peaks if event_start_sample < p < expected_event_end_sample]
+    
+    # If we don't have enough peaks, try lowering the threshold
+    if len(relevant_peaks) < n_jumpers - 1:  # We need n_jumpers-1 transitions
+        for lower_height in [0.7, 0.6, 0.5, 0.4, 0.3]:
+            peaks, _ = find_peaks(corr, height=lower_height, distance=fs)
+            relevant_peaks = [p for p in peaks if event_start_sample < p < expected_event_end_sample]
+            if len(relevant_peaks) >= n_jumpers - 1:
+                break
+    
+    # If we still don't have enough peaks, we'll need to estimate some transitions
+    relay_length_frames = fps * relay_length
+    
+    # Process peaks to identify jumper transitions
+    if len(relevant_peaks) >= n_jumpers - 1:
+        # Ideal case: we found enough beeps for transitions
+        # Sort peaks by time to ensure correct order
+        relevant_peaks.sort()
+        
+        # Use the first n_jumpers-1 peaks as transition points
+        transition_frames = [int(p / fs * fps) for p in relevant_peaks[:n_jumpers-1]]
+        
+        # Set ends for jumpers based on transition points
+        for i in range(n_jumpers - 1):
+            ends.append(transition_frames[i])
+            starts.append(transition_frames[i])
+        
+        # Add end for the last jumper
+        ends.append(expected_event_end)
+    else:
+        # Not enough peaks detected, use expected relay_length to estimate
+        for i in range(n_jumpers):
+            if i == 0:
+                # First jumper starts at event_start (already added to starts)
+                jumper_end = event_start + relay_length_frames
+                ends.append(jumper_end)
+                if i < n_jumpers - 1:
+                    starts.append(jumper_end)
+            elif i < n_jumpers - 1:
+                jumper_end = starts[i] + relay_length_frames
+                ends.append(jumper_end)
+                starts.append(jumper_end)
+            else:
+                # Last jumper
+                jumper_end = starts[i] + relay_length_frames
+                ends.append(jumper_end)
+    
+    # Validate and adjust if necessary
+    # Make sure all intervals are close to relay_length
+    for i in range(n_jumpers):
+        interval = ends[i] - starts[i]
+        # If an interval is significantly different from relay_length, adjust it
+        if abs(interval - relay_length_frames) > relay_length_frames * 0.2:  # 20% tolerance
+            # Adjust the end time to match expected relay_length
+            ends[i] = starts[i] + relay_length_frames
+            # If not the last jumper, adjust the next start time
+            if i < n_jumpers - 1:
+                starts[i + 1] = ends[i]
+    
+    # Final check: ensure the total length matches expected
+    if ends[-1] != expected_event_end:
+        # Adjust the last end to match the expected total event end
+        ends[-1] = expected_event_end
+    
+    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,
+              count_only_api, api_key,
               img_size=256, 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)
     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)
@@ -138,6 +264,10 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         event_length = int(event_length)
         event_start, event_end = detect_beeps(in_video, event_length)
         print(event_start, event_end)
+        if relay_detection_on:
+            n_jumpers = int(int(event_length) / int(relay_length))
+            relay_starts, relay_ends = detect_relay_beeps(in_video, event_start, int(relay_length), n_jumpers)
+            print(relay_starts, relay_ends)
     
     cap = cv2.VideoCapture(in_video)
     length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -184,7 +314,7 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     batch_list = []
     idx_list = []
     inference_futures = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
         for i in range(0, length + stride_length - stride_pad, stride_length):
             batch = all_frames[i:i + seq_len]
             Xlist = []
@@ -233,6 +363,7 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                 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
+            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]
     periodicity = np.divide(periodicities, period_length_overlaps, where=period_length_overlaps!=0)[:length]
@@ -254,6 +385,12 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
             print(f"Event detected: {event_start} - {event_end}")
             periodicity[:event_start] = 0
             periodicity[event_end:] = 0
+        if relay_detection_on:
+            for start, end in zip(relay_starts, relay_ends):
+                if start > 0 and end > 0:
+                    print(f"Relay Event detected: {start} - {end}")
+                    # immediately after the beep set periodicity to 0 for switch_delay seconds
+                    periodicity[start:start + int(float(switch_delay) * fps)] = 0
     pred_marks_peaks, _ = find_peaks(full_marks, distance=3, height=marks_threshold)
     full_marks_mask = np.zeros(len(full_marks))
     full_marks_mask[pred_marks_peaks] = 1
@@ -287,6 +424,11 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
             confidence = 0
         else:
             confidence = (np.mean(periodicity[periodicity > miss_threshold]) - miss_threshold) / (1 - miss_threshold)
+        periodicity_mask = periodicity > miss_threshold
+        if np.sum(periodicity_mask) == 0:
+            confidence = 0
+        else:
+            confidence = (np.mean(periodicity[periodicity > miss_threshold]) - miss_threshold) / (1 - miss_threshold)
     except ZeroDivisionError:
         confidence = 0
     self_err = abs(count_pred - marks_count_pred)
@@ -300,8 +442,16 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         count_msg = f"## Reps 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"## Reps Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
 
     if api_call:
+        if CACHE_API_CALLS:
+            # write outputs as row of csv
+            with open('api_calls.tsv', 'a') as f:
+                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")
         if CACHE_API_CALLS:
             # write outputs as row of csv
             with open('api_calls.tsv', 'a') as f:
@@ -359,6 +509,12 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     if beep_detection_on:
         # add vertical lines for beep event
         fig.add_vrect(x0=event_start / fps, x1=event_end / fps, fillcolor="LightSalmon", opacity=0.25, layer="below", line_width=0)
+        if relay_detection_on:
+            for start, end in zip(relay_starts, relay_ends):
+                start += 10  # add some padding
+                end -= 10
+                fig.add_vrect(x0=start / fps, x1=end / fps, fillcolor="LightGreen", opacity=0.25, layer="below",
+                              line_width=0)
 
     
     fig.update_layout(legend=dict(
@@ -405,10 +561,19 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     except FileNotFoundError:
         pass
     
+    try:
+        os.remove('temp.wav')
+    except FileNotFoundError:
+        pass
+    
     return in_video, count_msg, fig, hist, bar
         
 
 with gr.Blocks() as demo:
+    with gr.Row():
+        in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', 
+                                width=400, height=400, interactive=True, container=True,
+                                max_length=300)
     with gr.Row():
         in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', 
                                 width=400, height=400, interactive=True, container=True,
@@ -417,11 +582,16 @@ with gr.Blocks() as demo:
         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)
         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)
+            relay_detection_on = gr.Checkbox(label="Relay Event", elem_id='relay-beeps', visible=True)
+            relay_length = gr.Textbox(label="Relay Length (s)", elem_id='relay-length', visible=True, value='30')
+            switch_delay = gr.Textbox(label="Expected Switch Delay (s)", elem_id='event-length', visible=True, value='0.2')
         with gr.Column(min_width=480):
             out_video = gr.PlayableVideo(label="Video Clip", elem_id='output-video', format='mp4', width=400, height=400)
             
@@ -448,24 +618,33 @@ 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], 
+    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])
     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, count_only, api_token], 
+    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')
     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],
+        [None, 'https://hiemdall-dev2.azurewebsites.net/api/playlist/rec_UGEhqlMh/vod', '00:00:18', '00:00:55', True, 30],
         #['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], 
+                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],
                 fn=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
 
 
 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,