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 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")
c = os.path.join(os.path.dirname(__file__), "files", "train3.mp4")
d = os.path.join(os.path.dirname(__file__), "files", "train29.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, both_feet, has_misses, true_count, center_crop, img_resize, miss_threshold, img_size=192, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, median_pred_filter=True, api_call=False):
    print(x)
    img_size = int((img_size - 64) * img_resize + 64)
    api = HfApi(token=os.environ['DATASET_SECRET'])
    out_file = str(uuid.uuid1())
    if has_misses:
        out_file = "misses_" + out_file
    if true_count != -1:
        out_file += '_' + str(true_count)
        out_file = f"labeled_videos/{out_file}.mp4"
        api.upload_file(
            path_or_fileobj=x,
            path_in_repo=out_file,
            repo_id="dylanplummer/jumprope",
            repo_type="dataset",
        )
    
        
    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 api_call:
        return np.array2string(periodLength, formatter={'float_kind':lambda x: "%.3f" % x}).replace('\n', ''), np.array2string(periodicity, formatter={'float_kind':lambda x: "%.3f" % x}).replace('\n', '')
    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

    animate_frames = len(count)
    anim_subsample = max(1, int(fps / 24))
    fig, ax = plt.subplots(figsize = (3, 3))
    canvas_width, canvas_height = fig.canvas.get_width_height()
    img_ax = ax

    # imgs = []
    # for img in all_frames:
    #     img = Image.fromarray(img)
    #     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))))
    #     h_center, w_center = height / 2, width / 2
    #     h_start, w_start = int(h_center - img_size / 2), int(w_center - img_size / 2)
    #     cropped = img.crop((w_start, h_start, w_start + img_size, h_start + img_size))
    #     imgs.append(cropped)

    # confidence_analysis(periodicity, count, imgs)

    # alpha=1.0
    # colormap=plt.cm.OrRd
    # h, w, _ = np.shape(imgs[0])
    # wedge_x = 34 / canvas_width * w
    # wedge_y = 34 / canvas_height * h
    # wedge_r = 30 / canvas_height * h
    # txt_x = 34 / canvas_width * w
    # txt_y = 36 / canvas_height * h
    # otxt_size = 25 / canvas_height * h
    # wedge1 = matplotlib.patches.Wedge(
    #     center=(wedge_x, wedge_y),
    #     r=wedge_r,
    #     theta1=0,
    #     theta2=0,
    #     color=colormap(1.),
    #     alpha=alpha)
    # wedge2 = matplotlib.patches.Wedge(
    #     center=(wedge_x, wedge_y),
    #     r=wedge_r,
    #     theta1=0,
    #     theta2=0,
    #     color=colormap(0.5),
    #     alpha=alpha)

    # im = img_ax.imshow(cropped)

    # img_ax.add_patch(wedge1)
    # img_ax.add_patch(wedge2)
    # txt = img_ax.text(
    #     txt_x,
    #     txt_y,
    #     '0',
    #     size=otxt_size,
    #     ha='center',
    #     va='center',
    #     alpha=0.9,
    #     color='white',
    # )

    # def animate_fn(i):
    #     if anim_subsample:
    #         i *= anim_subsample
    #     cropped = imgs[i + stride_pad]
    #     current_count = count[i]
    #     if current_count % 2 == 0:
    #         wedge1.set_color(colormap(1.0))
    #         wedge2.set_color(colormap(0.5))
    #     else:
    #         wedge1.set_color(colormap(0.5))
    #         wedge2.set_color(colormap(1.0))
    #     txt.set_text(int(current_count))
    #     wedge1.set_theta1(-90)
    #     wedge1.set_theta2(-90 - 360 * (1 - current_count % 1.0))
    #     wedge2.set_theta1(-90 - 360 * (1 - current_count % 1.0))
    #     wedge2.set_theta2(-90)
        
    #     im.set_data(cropped)
    #     img_ax.set_title(f"Time: {i / fps:.1f}s, {current_count:.1f}/{count_pred:.1f} jumps")
    #     img_ax.set_xticks([])
    #     img_ax.set_yticks([])
    #     img_ax.spines['top'].set_visible(False)
    #     img_ax.spines['right'].set_visible(False)
    #     img_ax.spines['bottom'].set_visible(False)
    #     img_ax.spines['left'].set_visible(False)

    outf = x
    # anim_start_time = time.time()
    # # Open an ffmpeg process
    # outf = x.replace('.mp4', '_jump.mp4')
    # cmdstring = ('ffmpeg', 
    #             '-y', '-r', f'{30 if anim_subsample != 1 else int(fps)}', # overwrite, 24fps
    #             '-s', f'{canvas_width}x{canvas_height}',
    #             '-pix_fmt', 'argb',
    #             '-hide_banner', '-loglevel', 'error',
    #             '-f', 'rawvideo',  '-i', '-', # tell ffmpeg to expect raw video from the pipe
    #             '-i', x, '-map', '0:v', '-map', '1:a', # map video from the pipe, audio from the input file
    #             '-c:v', 'libx264', # https://trac.ffmpeg.org/wiki/Encode/H.264
    #             outf) # output encoding
    # try:
    #     p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE)
    # except FileNotFoundError as e:
    #     print(e)
    #     print('Trying to install ffmpeg...')
    #     os.system("apt install ffmpeg -y")
    #     p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE)

    # # Draw frames and write to the pipe
    # anim_length = int(animate_frames / anim_subsample) - 1
    # for frame in range(anim_length):
    #     # draw the frame
    #     animate_fn(frame)
    #     fig.canvas.draw()
    #     # extract the image as an ARGB string
    #     string = fig.canvas.tostring_argb()
    #     # write to pipe
    #     p.stdin.write(string)

    # # Finish up
    # p.communicate()
    # print(f"Animation done in {time.time() - anim_start_time:.2f} seconds")
    #cvrt_string = f'ffmpeg -hide_banner -loglevel error -i "{outf}" -i "{x}" -map 0:v -map 1:a -y -r 24 -s {canvas_width}x{canvas_height} -c:v libx264 "{outf.replace(".mp4", "_audio.mp4")}"'
    #os.system(cvrt_string)

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

    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])
    
    vid = px.imshow(np.uint8(all_frames)[:128], animation_frame=0, binary_string=True, binary_compression_level=5, binary_format='jpg', template="plotly_dark")
    vid.update_xaxes(showticklabels=False).update_yaxes(showticklabels=False)
    vid.layout.updatemenus[0].buttons[0].args[1]['frame']['duration'] = 13
    vid.layout.updatemenus[0].buttons[0].args[1]['transition']['duration'] = 5

    return count_msg, vid, 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(type="file", label="Input Video", elem_id='input-video', format='mp4').style(width=400)
            with gr.Column():
                gr.Markdown(label="Optional settings and parameters:")
                true_count = gr.Number(label="True Count (optional)", info="Provide a true count if you are ok with us using your video for training", elem_id='true-count', value=-1)
                both_feet = gr.Checkbox(label="Both feet", info="Count both feet rather than only one", elem_id='both-feet', value=False)
                misses = gr.Checkbox(label="Contains misses", info="Only necessary if providing a true count for training", elem_id='both-feet', value=False)
                center_crop = gr.Checkbox(label="Center crop square", info="Either crop a square out of the center or stretch to a square", elem_id='center-crop', value=True)
                image_size = gr.Slider(label="Image size", info="Lower image size is faster but less accurate", elem_id='miss-thresh', minimum=0.0, maximum=1.0, step=0.01, value=1.0)
                miss_thresh = gr.Slider(label="Miss threshold", info="Lower values are more sensitive to misses", elem_id='miss-thresh', minimum=0.0, maximum=0.99, step=0.01, value=0.95)
                with gr.Row():
                    run_button = gr.Button(label="Run", elem_id='run-button').style(full_width=False)
                    count_only = gr.Button(label="Run (No Viz)", elem_id='count-only', visible=False).style(full_width=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_video = gr.Plot(label="Output Video", elem_id='output-video')
        out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
        out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
                
    inputs = [in_video, both_feet, misses, true_count, center_crop, image_size, miss_thresh]
    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],
                        [c, False, False, -1, True, 1.0, 0.95],
                        [d, False, False, -1, True, 1.0, 0.95]
                    ],
                inputs=inputs,
                outputs=[out_text, out_video, out_plot, out_hist],
                fn=inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
    with gr.Accordion(label="Data usage and disclaimer", open=False):
        data_usage = "## Data usage:"
        data_usage += "\n* By default, no data submitted to this demo is stored by us."
        data_usage += "\n* If you would like to contribute your video for further model improvements please provide the true count (either one or both feet) and specify if the video contains any misses."
        data_usage += "\n* The video will be uploaded to a private dataset repository here on HuggingFace"
        gr.Markdown(data_usage)

        disclaimer = "## Disclaimer:"
        disclaimer += "\n* This model was trained on a small dataset of videos (~20 hours). It is not guaranteed to work on all videos and should not be used yet in real competitive settings."
        disclaimer += "\n* Deep learning models such as this one are susceptible to biases in the training data."
        disclaimer += " We are aware of the potential for bias and expect quantifiable differences in performance on counting videos from different demographics not represented in the training data. We are working to improve the model and mitigate these biases. If you notice any issues, please let us know."
        gr.Markdown(disclaimer)

    run_button.click(inference, inputs, outputs=[out_text, out_video, out_plot, out_hist])
    api_inference = partial(inference, api_call=True)
    count_only.click(api_inference, inputs, outputs=[period_length, periodicity], api_name='inference')


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