add progress and inference options

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import os
 import cv2
 import math
+import time
 import json
 import subprocess
 import matplotlib
@@ -29,14 +30,14 @@ from hls_download import download_clips
 
 #plt.style.use('dark_background')
 
-LOCAL = False
+LOCAL = True
 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'])
-
-#onnx_file = 'nextjump.onnx'
+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 = f'{current_model}.onnx'
+api = HfApi()
 
 if torch.cuda.is_available():
     print("Using CUDA")
@@ -80,44 +81,102 @@ def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
 
-def detect_beeps(video_path, event_length=30, beep_height=0.8):
+def detect_beeps(video_path, target_event_length=30, beep_height=0.8):
+    """
+    Detects beep sounds in a video file and returns frame indices for start and end points.
+    Finds the pair of peaks that are closest to the target event length.
+    
+    Args:
+        video_path: Path to the video file
+        target_event_length: Target duration of the event in seconds
+        beep_height: Initial threshold for peak detection
+        
+    Returns:
+        event_start: Frame index for the start of the event
+        event_end: Frame index for the end of the event
+    """
+    
+    # Read reference beep
     reference_file = 'beep.WAV'
     fs, beep = wavfile.read(reference_file)
     beep = beep[:, 0] + beep[:, 1]  # combine stereo to mono
+    
+    # Open video file
     video = cv2.VideoCapture(video_path)
+    length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = int(video.get(cv2.CAP_PROP_FPS))
+    
+    # Clean up any previous temporary files
     try:
         os.remove('temp.wav')
     except FileNotFoundError:
         pass
+    
+    # Extract audio from video
     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]
+    
+    # Read the extracted audio
+    _, audio = wavfile.read('temp.wav')
     audio = (audio[:, 0] + audio[:, 1]) / 2  # combine stereo to mono
+    
+    # Cross-correlate with the reference beep
     corr = correlate(audio, beep, mode='same') / audio.size
-    # min max scale to -1, 1
+    
+    # Min-max scale correlation to [-1, 1]
     corr = 2 * (corr - np.min(corr)) / (np.max(corr) - np.min(corr)) - 1
-    event_start = length
-    while length - event_start < fps * event_length:
-        peaks, _ = find_peaks(corr, height=beep_height, distance=fs)
-        event_start = int(peaks[0] / fs * fps)
-        event_end = int(peaks[-1] / fs * fps)
-        if event_end == event_start:
-            event_end = event_start + fps * event_length
-        beep_height -= 0.1
-        if beep_height <= 0.1:
-            event_start = 0
-            event_end = length
+    
+    # Target number of frames for the event
+    target_frames = fps * target_event_length
+    
+    # Strategy: Try different height thresholds to find peaks,
+    # then select the pair closest to the target length
+    best_pair = None
+    best_diff = float('inf')
+    
+    min_height = 0.3  # Minimum threshold to consider
+    height_step = 0.05  # Decrease step
+    
+    # Try different height thresholds
+    current_height = beep_height
+    while current_height >= min_height:
+        peaks, _ = find_peaks(corr, height=current_height, distance=fs//2)
+        
+        if len(peaks) >= 2:
+            # Check all possible pairs of peaks
+            for i in range(len(peaks)):
+                for j in range(i+1, len(peaks)):
+                    start_frame = int(peaks[i] / fs * fps)
+                    end_frame = int(peaks[j] / fs * fps)
+                    duration = end_frame - start_frame
+                    
+                    # Calculate how close this pair is to the target length
+                    diff = abs(duration - target_frames)
+                    
+                    # Update if this is the best match so far
+                    if diff < best_diff:
+                        best_diff = diff
+                        best_pair = (start_frame, end_frame)
+        if best_diff < 15: # If we found a good pair, break early
             break
-    #peaks, _ = find_peaks(corr, height=0.7, distance=fs)
-    #event_start = int(peaks[0] / fs * fps)
-    #event_end = int(peaks[-1] / fs * fps)
-    # plt.plot(corr)
-    # plt.plot(peaks, corr[peaks], "x")
-    # plt.savefig('beep.png')
-    # plt.close()
-
+        
+        # Reduce height threshold and try again
+        current_height -= height_step
+    
+    # If we found a good pair, use it
+    if best_pair:
+        event_start, event_end = best_pair
+    else:
+        # Fallback: use the whole video
+        event_start = 0
+        event_end = length
+    
+    # Optional visualization (commented out)
+    plt.plot(corr)
+    plt.plot(peaks, corr[peaks], "x")
+    plt.savefig('beep.png')
+    plt.close()
+    
     return event_start, event_end
 
 
@@ -236,16 +295,47 @@ 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 upload_video(out_text, in_video):
+    if out_text != '':
+        # generate a timestamp name for the video
+        upload_path = f"{int(time.time())}.mp4"
+        api.upload_file(
+            path_or_fileobj=in_video,
+            path_in_repo=upload_path,
+            repo_id="lumos-motion/single-rope-contest",
+            repo_type="dataset",
+        )
+
+
+def inference(in_video, stream_url, start_time, end_time, 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...")
+    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)})
     if in_video is None:
-        in_video = download_clips(stream_url, os.path.join(os.getcwd(), 'clips'), start_time, end_time)
+        in_video = download_clips(stream_url, os.path.join(os.getcwd(), 'clips'), start_time, end_time, use_60fps=use_60fps)
     else:  # local uploaded video (still resize with ffmpeg)
-        in_video = download_clips(in_video, os.path.join(os.getcwd(), 'clips'), start_time, end_time)
+        in_video = download_clips(in_video, os.path.join(os.getcwd(), 'clips'), start_time, end_time, use_60fps=use_60fps)
     progress(0, desc="Running inference...")
     has_access = False
     if api_call:
@@ -283,14 +373,15 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
         frame = cv2.cvtColor(np.uint8(frame), cv2.COLOR_BGR2RGB)
         # add square padding with opencv
         #frame = square_pad_opencv(frame)
-        frame_center_x = frame.shape[1] // 2
-        frame_center_y = frame.shape[0] // 2
-        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]
+        # frame_center_x = frame.shape[1] // 2
+        # frame_center_y = frame.shape[0] // 2
+        # 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]
+        frame = cv2.resize(frame, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC)
         all_frames.append(frame)
         
     cap.release()
@@ -310,8 +401,8 @@ 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=2) as executor:
-        for i in range(0, length + stride_length - stride_pad, stride_length):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        for i in progress.tqdm(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)]
@@ -339,25 +430,31 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
                 idx_list.append(idx_list[-1])
             future = executor.submit(run_inference, batch_list)
             inference_futures.append((batch_list, idx_list, future))
-        
+        progress(0, desc="Processing results...")
         # Collect and process the inference results
-        for batch_list, idx_list, future in tqdm(inference_futures):
+        for batch_list, idx_list, future in progress.tqdm(tqdm(inference_futures)):
             outputs = future.result()
             y1_out = outputs[0]
             y2_out = outputs[1]
             y3_out = outputs[2]
             y4_out = outputs[3]
             y5_out = outputs[4]
-            y6_out = outputs[5]
+            try:
+                y6_out = outputs[5]
+            except IndexError:
+                y6_out = np.zeros((len(batch_list), seq_len, 2))
             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]
+                try:
+                    period_lengths_rope[idx:idx+seq_len] += periodLength[:, 1]
+                except IndexError:
+                    period_lengths_rope[idx:idx+seq_len] += periodLength[:, 0]
                 periodicities[idx:idx+seq_len] += periodicity
                 full_marks[idx:idx+seq_len] += marks
                 event_type_logits[idx:idx+seq_len] += event_type
@@ -404,16 +501,24 @@ def inference(in_video, stream_url, start_time, end_time, beep_detection_on, eve
     periodicity_mask = np.int32(periodicity > miss_threshold)
     numofReps = 0
     count = []
+    miss_detected = True
+    num_misses = 0
     for i in range(len(periodLength)):
         if periodLength[i] < 2 or periodicity_mask[i] == 0:
             numofReps += 0
+            if not miss_detected:
+                miss_detected = True
+                num_misses += 1
+                numofReps -= 2
         elif full_marks_mask[i]:  # high confidence mark detected
             if math.modf(numofReps)[0] < 0.2:  # probably false positive/late detection
                 numofReps = float(int(numofReps))
             else:
                 numofReps = float(int(numofReps) + 1.01)  # round up
+            miss_detected = False
         else:
             numofReps += max(0, periodicity_mask[i]/(periodLength[i]))
+            miss_detected = False
         count.append(round(float(numofReps), 2))
     count_pred = count[-1]
     marks_count_pred = 0
@@ -697,13 +802,36 @@ with gr.Blocks() as demo:
                                 max_length=300)
     with gr.Row():
         with gr.Column():
+            gr.Markdown(
+                """
+                ### Stream Input Options
+                Either upload a video file above, or provide a stream URL below.
+                """,
+                elem_id='stream-input-options',
+            )
             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():
+            gr.Markdown(
+                """
+                ### Inference Options
+                Select the model and framerate for inference.
+                """,
+                elem_id='inference-options',
+            )
+            use_60fps = gr.Checkbox(label="Use 60 FPS", elem_id='use-60fps', visible=True)
+            model_choice =  gr.Dropdown(
+                ["nextjump_speed", "nextjump_all", "nextjump_both_feet"], label="Model Choice", info="For now just speed-only or general model",
+            )
+        with gr.Column():
+            gr.Markdown(
+                """
+                ### Beep Detection Options
+                Must be using official IJRU timing tracks.
+                """,
+                elem_id='beep-detection-options',
+            )
             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)
@@ -740,19 +868,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, in_stream_url, in_stream_start, in_stream_end, 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, in_stream_url, in_stream_start, in_stream_end, 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', '', '00:00:00', '', 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, in_stream_url, in_stream_start, in_stream_end, 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",