app.py

import gradio as gr
import numpy as np
from PIL import Image
import os
import cv2
import math
import time
import json
import subprocess
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import concurrent.futures
from scipy.io import wavfile
from scipy.signal import medfilt, correlate, find_peaks
from functools import partial
from passlib.hash import pbkdf2_sha256
from tqdm import tqdm
import pandas as pd
import plotly.express as px
import onnxruntime as ort
import torch
from torchvision import transforms
import torchvision.transforms.functional as F

from huggingface_hub import hf_hub_download
from huggingface_hub import HfApi

from hls_download import download_clips

#plt.style.use('dark_background')

LOCAL = False
IMG_SIZE = 256
CACHE_API_CALLS = False
os.makedirs(os.path.join(os.getcwd(), 'clips'), exist_ok=True)
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")
    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)})

    
def square_pad_opencv(image):
    h, w = image.shape[:2]
    max_wh = max(w, h)
    hp = int((max_wh - w) / 2)
    vp = int((max_wh - h) / 2)
    return cv2.copyMakeBorder(image, vp, vp, hp, hp, cv2.BORDER_CONSTANT, value=[0, 0, 0])


def preprocess_image(img, img_size):
    #img = square_pad_opencv(img)
    #img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_CUBIC)
    img = Image.fromarray(img)
    transforms_list = []
    transforms_list.append(transforms.ToTensor())
    preprocess = transforms.Compose(transforms_list)
    return preprocess(img).unsqueeze(0)


def run_inference(batch_X):
    global ort_sess
    batch_X = torch.cat(batch_X)
    return ort_sess.run(None, {'video': batch_X.numpy()})


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


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'
    print(audio_convert_command)
    subprocess.call(audio_convert_command, shell=True)
    
    # 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 correlation to [-1, 1]
    corr = 2 * (corr - np.min(corr)) / (np.max(corr) - np.min(corr)) - 1
    
    # 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
        
        # 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


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 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 count_phases(phase_sin, phase_cos, threshold=0.5):
    """
    Count the number of phase transitions in the sine and cosine phases.
    
    Args:
        phase_sin: Numpy array of sine phase values
        phase_cos: Numpy array of cosine phase values
        threshold: Threshold to consider a transition
    Returns:
        count: Number of phase transitions
        phase_indices: Indices where transitions occur
    """
    phase_indices = []
    count = 0
    for i in range(1, len(phase_sin)):
        # Check if the sine and cosine phases cross each other
        if (phase_sin[i-1] < threshold and phase_sin[i] >= threshold) or \
           (phase_sin[i-1] >= threshold and phase_sin[i] < threshold):
            # Check if the cosine phase crosses the threshold
            if (phase_cos[i-1] < threshold and phase_cos[i] >= threshold) or \
            (phase_cos[i-1] >= threshold and phase_cos[i] < threshold):
                phase_indices.append(i)
                count += 1
    return count, phase_indices



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=2,
              miss_threshold=0.5, marks_threshold=0.5, median_pred_filter=True, both_feet=True, 
              api_call=False,
              progress=gr.Progress()):
    global current_model, ort_sess
    print(in_video)
    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'])
        #onnx_file = f'{current_model}.onnx'

        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:
        print("No video input provided.")
        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)
        print("Using uploaded video input.")
        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:
        has_access = pbkdf2_sha256.verify(os.environ['DEV_API_TOKEN'], api_key)
        if not has_access:
            return "Invalid API Key"
    
    if beep_detection_on:
        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))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    period_length_overlaps = np.zeros(length + seq_len)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    seconds = length / fps
    all_frames = []
    frame_i = 0
    resize_amount = max((IMG_SIZE + 64) / frame_width, (IMG_SIZE + 64) / frame_height)
    while cap.isOpened():
        frame_i += 1
        
        ret, frame = cap.read()
        if ret is False:
            frame = all_frames[-1]  # padding will be with last frame
            break
        
        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 = cv2.resize(frame, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC)
        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
        all_frames.append(all_frames[-1])
    batch_list = []
    idx_list = []
    inference_futures = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) 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)]
            for idx, future in sorted(preprocess_tasks, key=lambda x: x[0]):
                Xlist.append(future.result())
            
            if len(Xlist) < seq_len:
                for _ in range(seq_len - len(Xlist)):
                    Xlist.append(Xlist[-1])
            
            X = torch.cat(Xlist)
            X *= 255
            batch_list.append(X.unsqueeze(0))
            idx_list.append(i)
            
            if len(batch_list) == batch_size:
                future = executor.submit(run_inference, batch_list)
                inference_futures.append((batch_list, idx_list, future))
                batch_list = []
                idx_list = []
        # Process any remaining batches
        if batch_list:
            while len(batch_list) != batch_size:
                batch_list.append(batch_list[-1])
                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 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]
            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
                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]
                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
                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))
    per_frame_event_types = np.argmax(per_frame_event_type_logits, axis=1)
    
    if median_pred_filter:
        periodicity = medfilt(periodicity, 5)
        periodLength = medfilt(periodLength, 5)
    periodicity = sigmoid(periodicity)
    full_marks = sigmoid(full_marks)
    # if the event_start and event_end (in frames) are detected and form a valid event of event_length (in seconds)
    if beep_detection_on:
        if event_start > 0 and event_end > 0 and (event_end - event_start) - (event_length * fps) < 0.5:
            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
    periodicity_mask = np.int32(periodicity > miss_threshold)
    phase_count, phase_indices = count_phases(phase_sin, phase_cos, threshold=-0.5)
    numofReps = 0
    count = []
    miss_detected = True
    num_misses = -1  # end of event is not counted as a miss
    miss_frames = []
    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
                miss_frames.append(i)
                #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
    for i in range(len(full_marks) - 1):
        # if a jump was counted, and periodicity is high, and the next frame was not counted (to avoid double counting)
        if full_marks_mask[i] > 0 and periodicity_mask[i] > 0 and full_marks_mask[i + 1] == 0:
            marks_count_pred += 1
    
    if not both_feet:
        count_pred = count_pred / 2
        marks_count_pred = marks_count_pred / 2
        count = np.array(count) / 2
    try:
        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)
    try:
        self_pct_err = self_err / count_pred
    except ZeroDivisionError:
        self_pct_err = 0
    total_confidence = confidence * (1 - self_pct_err)

    # find the fastest second (30 frames if 30fp and 60 frames if 60fps) based on the period_length
    scan_window = 60 if use_60fps else 30
    fastest_frames_start = 0
    fastest_period = float('inf')
    for i in range(0, len(periodLength) - scan_window, scan_window // 2):
        #if np.sum(periodicity_mask[i:i + scan_window]) > 0:
        avg_period = np.mean(periodLength[i:i + scan_window])
        if avg_period < fastest_period:
            fastest_period = avg_period
            fastest_frames_start = i
    fastest_frames_end = fastest_frames_start + scan_window
    fastest_jumps_per_second = np.clip(1 / ((fastest_period / fps) + 0.0001), 0, 10)
    print(f"Fastest jumps per second: {fastest_jumps_per_second:.2f} (from frames {fastest_frames_start} to {fastest_frames_end})")

    # measure the reaction time to the beep (if beep detection is on) as the time to reach average speed 
    time_to_speed = 0
    if beep_detection_on:
        avg_speed = np.mean(periodLength[periodicity_mask])
        reaction_frame = np.argmax((periodLength < avg_speed) & (periodicity_mask))
        print(f"Reaction frame: {reaction_frame}, Avg Speed: {avg_speed}")
        time_to_speed = (reaction_frame - event_start) / fps

    # get peak speed and lowest speed
    peak_speed = np.quantile(periodLength[periodicity_mask], 0.01) if np.any(periodicity_mask) else 0
    lowest_speed = np.quantile(periodLength[periodicity_mask], 0.99) if np.any(periodicity_mask) else 0
    peak_jps = np.clip(1 / ((peak_speed / fps) + 0.0001), 0, 10)
    lowest_jps = np.clip(1 / ((lowest_speed / fps) + 0.0001), 0, 10)
    slowdown = (lowest_jps - peak_jps)
    slowdown_percent = (slowdown / peak_jps) * 100 if peak_jps > 0 else 0

    print('slowdown', slowdown)
    print('percent', slowdown_percent)

    # estimate the score assuming no misses and fill in the gaps
    estimated_score = 0
    filled_periodLength = np.zeros(len(periodLength))
    started = False
    for i in range(len(periodLength)):
        if beep_detection_on and i < event_start:
            filled_periodLength[i] = 0
        elif beep_detection_on and i >= event_end:
            filled_periodLength[i] = 0
        elif periodicity_mask[i] > 0:
            started = True
            filled_periodLength[i] = periodLength[i]
        elif not started:
            filled_periodLength[i] = 0
        else:
            # fill in the gaps with the previous value
            filled_periodLength[i] = filled_periodLength[i - 1]
    estimated_score = 0
    for i in range(len(filled_periodLength)):
        if filled_periodLength[i] < 2:
            estimated_score += 0
        else:
            estimated_score += max(0, periodicity_mask[i] / (filled_periodLength[i]))
    print(f"Estimated score: {estimated_score:.2f}")
    
    # find the recovery times after each miss
    recovery_times = []
    if len(miss_frames) > 0:
        avg_speed = np.mean(periodLength[periodicity_mask])
        for miss_frame in miss_frames:
            # find the next frame where the speed is above avg_speed
            recovery_frame = np.argmax((periodLength[miss_frame:] > avg_speed) & (periodicity_mask[miss_frame:])) + miss_frame
            if recovery_frame > miss_frame:
                recovery_time = (recovery_frame - miss_frame) / fps
                recovery_times.append(recovery_time)
            else: # end of event
                pass 
    print(f"Recovery times: {recovery_times}")
            

    if LOCAL:
        if both_feet:
            count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks: {marks_count_pred:.1f}, Phase: {phase_count:.1f}, Confidence: {total_confidence:.2f}, Time to Speed: {time_to_speed:.2f} seconds"
        else:
            count_msg = f"## Reps Count (one foot): {count_pred:.1f}, Marks: {marks_count_pred:.1f}, Phase: {phase_count:.1f}, Confidence: {total_confidence:.2f}, Time to Speed: {time_to_speed:.2f} seconds"
    else:
        if both_feet:
            count_msg = f"## Reps 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}"

    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 count_only_api:
            return f"{count_pred:.2f} (conf: {total_confidence:.2f})"
        else:
            # 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}",
                "phase_count": f"{phase_count:.1f}",
                "confidence": f"{total_confidence:.2f}",
                "fastest_frames_start": fastest_frames_start,
                "fastest_frames_end": fastest_frames_end,
                "fastest_jumps_per_second": f"{fastest_jumps_per_second:.2f}",
                "lowest_jumps_per_second": f"{lowest_jps:.2f}",
                "fastest_period_length": f"{fastest_period:.2f}",
                "lowest_period_length": f"{lowest_speed:.2f}",
                "time_to_speed": f"{time_to_speed:.2f}" if beep_detection_on else 0,
                "slowdown": f"{slowdown:.2f}",
                "slowdown_percent": f"{slowdown_percent:.2f}",
                "num_misses": num_misses,
                "miss_frames": np.array2string(np.array(miss_frames[:num_misses]), formatter={'int':lambda x: str(x)}, threshold=np.inf).replace('\n', ''),
                "recovery_times": np.array2string(np.array(recovery_times), formatter={'float_kind':lambda x: "%.2f" % x}, threshold=np.inf).replace('\n', ''),
                "no_miss_score": f"{estimated_score:.2f}" if num_misses > 0 else f"{count_pred:.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(5, 1, figsize=(14, 10)) # Added a plot for count

    # # 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[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[4].plot(count)
    # axs[4].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
    jumps_per_second[misses] = 0
    frame_type = np.array(['miss' if miss else 'frame' for miss in misses])
    frame_type[full_marks > marks_threshold] = 'jump'
    per_frame_event_types = np.clip(per_frame_event_types, 0, 6) / 6
    df = pd.DataFrame.from_dict({'period length': periodLength, 
                                 '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,
                                'jumps': full_marks,
                                'jumps_size': (full_marks + 0.05) * 10,
                                'miss_size': np.clip((1 - periodicity) * 0.9 + 0.1, 1, 8),
                                'seconds': np.linspace(0, seconds, num=len(periodLength))})
    event_type_tick_vals = np.linspace(0, 1, num=7)
    event_type_colors = ['red', 'orange', 'green', 'blue', 'purple', 'pink', 'black']
    fig = px.scatter(data_frame=df,
                    x='seconds', 
                    y='jumps per second',
                    #symbol='frame_type',
                    #symbol_map={'frame': 'circle', 'miss': 'circle-open', 'jump': 'triangle-down'},
                    color='event_type',
                    size='jumps_size',
                    size_max=8,
                    color_continuous_scale=[(t, c) for t, c in zip(event_type_tick_vals, event_type_colors)],
                    range_color=(0,1),
                    title="Jumping speed (jumps-per-second)",
                    trendline='rolling',
                    trendline_options=dict(window=16),
                    trendline_color_override="goldenrod",
                    trendline_scope='overall',
                    template="plotly_dark")
    
    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(
            orientation="h",
            yanchor="bottom",
            y=0.98,
            xanchor="right",
            x=1,
            font=dict(
                family="Courier",
                size=12,
                color="black"
                ),
            bgcolor="AliceBlue",
        ),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)'
    )
    # 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'
    ))


    # -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
    phase_jumps = np.zeros(len(phase_sin))
    phase_jumps[phase_indices] = 1
    fig_phase_spiral = px.scatter(x=phase_cos, y=phase_sin,
                                color=phase_jumps,
                                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="Phase Jumps",))
    # 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", 
                        marginal="box",
                        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])
    try:
        os.remove('temp.wav')
    except FileNotFoundError:
        pass
    
    return 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', 
                                width=400, height=400, interactive=True, container=True,
                                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_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",
                value="nextjump_speed", elem_id='model-choice'
            )
        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)
            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.Row():
        run_button = gr.Button(value="Run", elem_id='run-button', scale=1)
        api_dummy_button = gr.Button(value="Run (No Viz)", elem_id='count-only', visible=False, scale=2)
        count_only = gr.Checkbox(label="Count Only", visible=False)
        api_token = gr.Textbox(label="API Key", elem_id='api-token', visible=False)

    with gr.Column(elem_id='output-video-container'):
        with gr.Row():
            with gr.Column():
                out_text = gr.Markdown(label="Predicted Count", elem_id='output-text')
                period_length = gr.Textbox(label="Period Length", elem_id='period-length', visible=False)
                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')
            with gr.Column():
                out_event_type_dist = gr.Plot(label="Event Type Distribution", elem_id='output-event-type-dist')
              

    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, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
                     outputs=[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, 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, '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, use_60fps, model_choice, beep_detection_on, event_length, relay_detection_on, relay_length, switch_delay], 
                outputs=[out_text, out_plot, out_phase_spiral, out_phase, out_hist, out_event_type_dist],
                fn=demo_inference, cache_examples=False)


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=True)
    else:
        demo.queue(api_open=True, max_size=15).launch(share=False)