| 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'])
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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])
|
|
|
| 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)
|
|
|
|
|
| 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]
|
| 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()
|
|
|
|
|
| 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):
|
| all_frames.append(all_frames[-1])
|
| batch_list = []
|
| idx_list = []
|
| print(length, stride_length, stride_pad)
|
| 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(transforms.CenterCrop((img_size, img_size)))
|
| else:
|
| transforms_list.append(transforms.Resize((img_size, img_size)))
|
|
|
|
|
| transforms_list += [
|
| transforms.ToTensor()]
|
|
|
| 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:
|
| 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)
|
|
|
|
|
| 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)
|
|
|
| demo_inference = partial(inference, count_only_api=False, api_key=None)
|
|
|
| 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=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
|
|
|
| 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) |
