app.py

import gradio as gr
import numpy as np
from PIL import Image
from openvino.runtime import Core
import os
import cv2
import uuid
import time
import subprocess
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.signal import medfilt
from functools import partial
from passlib.hash import pbkdf2_sha256
from tqdm import tqdm
import pandas as pd
import plotly.express as px
import torch
from torchvision import transforms
import torchvision.transforms.functional as F

from huggingface_hub import hf_hub_download
from huggingface_hub import HfApi

plt.style.use('dark_background')

hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.bin", repo_type="model", token=os.environ['DATASET_SECRET'])
model_xml = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.xml", repo_type="model", token=os.environ['DATASET_SECRET'])
hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.mapping", repo_type="model", token=os.environ['DATASET_SECRET'])
#model_xml = "model_ir/model.xml"

ie = Core()
model_ir = ie.read_model(model=model_xml)
config = {"PERFORMANCE_HINT": "LATENCY"}
compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)

a = os.path.join(os.path.dirname(__file__), "files", "dylan.mp4")
b = os.path.join(os.path.dirname(__file__), "files", "train14.mp4")

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


def confidence_analysis(periodicity, counts, frames, out_dir='confidence_animations', top_n=9):
    os.makedirs(out_dir, exist_ok=True)
    jump_arrs = []
    confidence_arrs = []
    current_jump = []
    current_confidence = []
    current_period = 1
    for i in range(len(periodicity)):
        if counts[i] < current_period:
            current_jump.append(np.array(frames[i]))
            current_confidence.append(periodicity[i])
        else:
            jump_arrs.append(current_jump)
            confidence_arrs.append(current_confidence)
            current_jump = [np.array(frames[i])]
            current_confidence = [periodicity[i]]
            current_period += 1
    avg_confidences = [np.median(x) for x in confidence_arrs]
    conf_order = np.argsort(avg_confidences)
    tiled_img = []
    tiled_confs = []
    for out_i, conf_idx in enumerate(conf_order):
        frames = np.array(jump_arrs[conf_idx])
        confidence = np.array(confidence_arrs[conf_idx])
        mean_confidence = np.median(confidence)
        tiled_img.append(frames)
        tiled_confs.append(mean_confidence)
        # fig, axs = plt.subplots(1, 1, figsize = (3, 3))

        # img_ax = axs
        # img_ax.imshow(np.zeros((128, 128, 3)))

        # def animate(i):
        #     img = frames[i]
        #     #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) 
        #     print(np.min(img), np.mean(img), np.max(img))
        #     img_ax.imshow(np.clip(img, 0, 255))
        #     print(i, end=' ')
        #     img_ax.set_xticks([])
        #     img_ax.set_yticks([])
        #     img_ax.set_title(f'Confidence: {mean_confidence:.2f}')


        # anim = FuncAnimation(fig, animate, frames=len(frames), interval=200)
        # anim.save(f'{out_dir}/conf_{out_i}.gif', writer=None)
        # plt.close(fig)
        if top_n > 10:
            break
    longest_len = max([len(x) for x in tiled_img])
    looped_tiled_img = []
    for frames in tiled_img:
        looped_tiled_img.append(np.concatenate([frames, frames[::-1]] * (longest_len // len(frames) + 1), axis=0)[:longest_len])
    # animate each tile
    n_rows = int(np.ceil(np.sqrt(top_n)))
    n_cols = int(np.ceil(top_n / n_rows))
    fig, axs = plt.subplots(n_rows, n_cols, figsize = (n_cols * 2, n_rows * 2))
    for i in range(n_rows):
        for j in range(n_cols):
            if i * n_cols + j < len(looped_tiled_img):
                img_ax = axs[i][j]
                img_ax.imshow(np.zeros((128, 128, 3)))
    def animate(i):
        print(i, end=' ')
        for row in range(n_rows):
            for col in range(n_cols):
                img = looped_tiled_img[row * n_cols + col][i]
                img_ax = axs[row][col]
                img_ax.imshow(np.clip(img, 0, 255))
                img_ax.set_xticks([])
                img_ax.set_yticks([])
                img_ax.set_title(f'Conf: {tiled_confs[row * n_cols + col]:.2f}')
    anim = FuncAnimation(fig, animate, frames=longest_len, interval=200)
    anim.save(f'{out_dir}/tiled_conf.gif', writer=None)
    plt.close(fig)


def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, miss_threshold=0.85, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
    print(x)
    #api = HfApi(token=os.environ['DATASET_SECRET'])
    #out_file = str(uuid.uuid1())
    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"
        
    cap = cv2.VideoCapture(x)
    length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    period_length_overlaps = np.zeros(length + seq_len)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    seconds = length / fps
    all_frames = []
    frame_i = 1
    while cap.isOpened():
        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)
        img = Image.fromarray(frame)
        width, height = img.size
        if width > height:
            img = img.resize((int(width / (height / img_size)), img_size))
        else:
            img = img.resize((img_size, int(height / (width / img_size))))
        all_frames.append(np.uint8(img))
        frame_i += 1
    cap.release()

    # Get output layer
    output_layer_period_length = compiled_model_ir.output(0)
    output_layer_periodicity = compiled_model_ir.output(1)
    length = len(all_frames)
    period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
    periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
    period_length_overlaps = np.zeros(len(all_frames) + seq_len + stride_length)
    for _ in range(seq_len + stride_length):  # pad full sequence
        all_frames.append(all_frames[-1])
    batch_list = []
    idx_list = []
    for i in tqdm(range(0, length + stride_length - stride_pad, stride_length)):
        batch = all_frames[i:i + seq_len]
        Xlist = []
        for img in batch:
            transforms_list = []
            if center_crop:
                #transforms_list.append(SquarePad())
                transforms_list.append(transforms.CenterCrop((img_size, img_size)))
            else:
                transforms_list.append(transforms.Resize((img_size, img_size)))
            

            transforms_list += [
                transforms.ToTensor()]
                #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
            preprocess = transforms.Compose(transforms_list)
            frameTensor = preprocess(Image.fromarray(img)).unsqueeze(0)
            Xlist.append(frameTensor)

        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:
            batch_X = torch.cat(batch_list)
            result = compiled_model_ir(batch_X)
            y1pred = result[output_layer_period_length]
            y2pred = result[output_layer_periodicity]
            for y1, y2, idx in zip(y1pred, y2pred, idx_list):
                periodLength = y1.squeeze()
                periodicity = y2.squeeze()
                period_lengths[idx:idx+seq_len] += periodLength
                periodicities[idx:idx+seq_len] += periodicity
                period_length_overlaps[idx:idx+seq_len] += 1
            batch_list = []
            idx_list = []
    if len(batch_list) != 0:  # still some leftover frames
        while len(batch_list) != batch_size:
            batch_list.append(batch_list[-1])
            idx_list.append(idx_list[-1])
        batch_X = torch.cat(batch_list)
        result = compiled_model_ir(batch_X)
        y1pred = result[output_layer_period_length]
        y2pred = result[output_layer_periodicity]
        for y1, y2, idx in zip(y1pred, y2pred, idx_list):
            periodLength = y1.squeeze()
            periodicity = y2.squeeze()
            period_lengths[idx:idx+seq_len] += periodLength
            periodicities[idx:idx+seq_len] += periodicity
            period_length_overlaps[idx:idx+seq_len] += 1
            
    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]
    
    if median_pred_filter:
        periodicity = medfilt(periodicity, 5)
        periodLength = medfilt(periodLength, 5)
    periodicity = sigmoid(periodicity)
    periodicity_mask = np.int32(periodicity > miss_threshold)
    numofReps = 0
    count = []
    for i in range(len(periodLength)):
        if periodLength[i] < 2 or periodicity_mask[i] == 0:
            numofReps += 0
        else:
            numofReps += max(0, periodicity_mask[i]/(periodLength[i]))
        count.append(round(float(numofReps), 2))
    count_pred = count[-1]
    if not both_feet:
        count_pred = count_pred / 2
        count = np.array(count) / 2

    if both_feet:
        count_msg = f"## Predicted Count (both feet): {count_pred:.1f}"
    else:
        count_msg = f"## Predicted Count (one foot): {count_pred:.1f}"

    if api_call:
        if count_only_api:
            return f"{count_pred:.2f}"
        else:
            return np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
                np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
                np.array2string(count_pred, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', '')
   

    jumps_per_second = np.clip(1 / ((periodLength / fps) + 0.05), 0, 8)
    jumping_speed = np.copy(jumps_per_second)
    misses = periodicity < miss_threshold
    jumps_per_second[misses] = 0
    df = pd.DataFrame.from_dict({'period length': periodLength, 
                                 'jumping speed': jumping_speed,
                                'jumps per second': jumps_per_second,
                                'periodicity': periodicity,
                                'miss': misses,
                                'miss_size': np.clip((1 - periodicity) * 0.9 + 0.1, 1, 10),
                                'seconds': np.linspace(0, seconds, num=len(periodLength))})
    fig = px.scatter(data_frame=df,
                    x='seconds', 
                    y='jumps per second',
                    symbol='miss',
                    symbol_map={False: 'triangle-down', True: 'circle-open'},
                    color='periodicity',
                    size='miss_size',
                    size_max=8,
                    color_continuous_scale='RdYlGn',
                    title="Jumping speed (jumps-per-second)",
                    trendline='rolling',
                    trendline_options=dict(window=32),
                    trendline_color_override="goldenrod",
                    trendline_scope='overall',
                    template="plotly_dark")
    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)'
    )

    hist = px.histogram(df, 
                        x="jumps per second", 
                        template="plotly_dark", 
                        marginal="box",
                        histnorm='percent',
                        title="Distribution of jumping speed (jumps-per-second)",
                        range_x=[np.min(jumps_per_second[jumps_per_second > 0]) - 0.5, np.max(jumps_per_second) + 0.5])

    return count_msg, fig, hist, periodLength
        

DESCRIPTION = '# NextJump'
DESCRIPTION += '\n## AI Counting for Competitive Jump Rope'
DESCRIPTION += '\nDemo created by [Dylan Plummer](https://dylan-plummer.github.io/). Check out the [NextJump iOS app](https://apps.apple.com/us/app/nextjump-jump-rope-counter/id6451026115).'


with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Column():
        with gr.Row():
            in_video = gr.Video(label="Input Video", elem_id='input-video', format='mp4', width=400, scale=2)
            
    with gr.Row():
        run_button = gr.Button(label="Run", elem_id='run-button', style=dict(full_width=False), scale=1)
        api_dummy_button = gr.Button(label="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.Column(min_width=480):
                #out_video = gr.PlayableVideo(label="Output Video", elem_id='output-video', format='mp4')
        out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
        out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
              
    with gr.Accordion(label="Instructions and more information", open=False):
        instructions = "## Instructions:"
        instructions += "\n* Upload a video and click 'Run' to get a prediction of the number of jumps (either one foot, or both). This could take a couple minutes! The model is trained on single rope and double dutch speed, but try out any videos you want."
        instructions += "\n* If you know the true count, you can provide it and your video will be used later to improve the model."
        instructions += "\n\n## Tips (optional):"
        instructions += "\n* Trim the video to start and end of the event"
        instructions += "\n* Frame the jumper fully, in the center of the frame"
        instructions += "\n* Videos are automatically resized, so higher resolution will not help, but a closer framing of the jumper might help. Try cropping the video differently."
        instructions += "\n\n\nUnfortunately due to inference costs, right now we have to deploy a slower, less accurate version of the model but it still works surprisingly well in most cases. If you run into any issues let us know. If you would like to contribute to the project, please reach out!"
        gr.Markdown(instructions)

        faq = "## FAQ:"
        faq += "\n* **Q:** Does the model recognize misses?\n    * **A:** Yes, but if it fails, you can try tuning the miss threshold slider to make it more sensitive."
        faq += "\n* **Q:** Does the model recognize double dutch?\n    * **A:** Yes, but it is trained on a smaller set of double dutch videos, so it may not work perfectly."
        faq += "\n* **Q:** Does the model recognize double unders\n    * **A:** Yes, but it is trained on a smaller set of double under videos, so it may not work perfectly. It is also trained to count the rope, not the jumps so you will need to divide the count by 2."
        faq += "\n* **Q:** Does the model count both feet?\n    * **A:** Yes, but for convention we usually return the halved (one foot) score."
        gr.Markdown(faq)

    gr.Examples(examples=[
                        [a, True, False, -1, True, 1.0, 0.95],
                        [b, False, True, -1, True, 1.0, 0.95],
                    ],
                inputs=[in_video],
                outputs=[out_text, out_plot, out_hist],
                fn=inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
    demo_inference = partial(inference, count_only_api=False, api_key=None)
    run_button.click(demo_inference, [in_video], outputs=[out_text, out_plot, out_hist])
    api_inference = partial(inference, api_call=True)
    api_dummy_button.click(api_inference, [in_video, count_only, api_token], outputs=[period_length], api_name='inference')


if __name__ == "__main__":
    demo.queue(api_open=True, max_size=15).launch(share=False)