Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import os | |
| import cv2 | |
| import math | |
| import spaces | |
| import matplotlib | |
| matplotlib.use('Agg') | |
| import matplotlib.pyplot as plt | |
| from scipy.signal import medfilt, 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 | |
| plt.style.use('dark_background') | |
| onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="nextjump.onnx", 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) | |
| class SquarePad: | |
| # https://discuss.pytorch.org/t/how-to-resize-and-pad-in-a-torchvision-transforms-compose/71850/9 | |
| def __call__(self, image): | |
| w, h = image.size | |
| max_wh = max(w, h) | |
| hp = int((max_wh - w) / 2) | |
| vp = int((max_wh - h) / 2) | |
| padding = (hp, vp, hp, vp) | |
| return F.pad(image, padding, 0, 'constant') | |
| def sigmoid(x): | |
| return 1 / (1 + np.exp(-x)) | |
| def inference(x, count_only_api, api_key, | |
| img_size=288, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, | |
| miss_threshold=0.8, marks_threshold=0.5, median_pred_filter=True, center_crop=True, both_feet=True, | |
| api_call=False, | |
| progress=gr.Progress()): | |
| progress(0, desc="Starting...") | |
| # check if GPU is available | |
| if torch.cuda.is_available(): | |
| providers = [("CUDAExecutionProvider", {"device_id": torch.cuda.current_device(), | |
| "user_compute_stream": str(torch.cuda.current_stream().cuda_stream)})] | |
| sess_options = ort.SessionOptions() | |
| ort_sess = ort.InferenceSession(onnx_file, sess_options=sess_options, providers=providers) | |
| else: | |
| ort_sess = ort.InferenceSession(onnx_file) | |
| #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) | |
| all_frames.append(img) | |
| frame_i += 1 | |
| cap.release() | |
| 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) | |
| full_marks = np.zeros(len(all_frames) + seq_len + stride_length) | |
| event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 7)) | |
| 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 = [] | |
| for i in tqdm(range(0, length + stride_length - stride_pad, stride_length)): | |
| batch = all_frames[i:i + seq_len] | |
| Xlist = [] | |
| print('Preprocessing...') | |
| for img in batch: | |
| transforms_list = [] | |
| # if center_crop: | |
| # if width > height: | |
| # transforms_list.append(transforms.Resize((int(width / (height / img_size)), img_size))) | |
| # else: | |
| # transforms_list.append(transforms.Resize((img_size, int(height / (width / img_size))))) | |
| # transforms_list.append(transforms.CenterCrop((img_size, img_size))) | |
| # else: | |
| transforms_list.append(SquarePad()) | |
| transforms_list.append(transforms.Resize((img_size, img_size), interpolation=Image.BICUBIC)) | |
| 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(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) | |
| print('Running inference...') | |
| if len(batch_list) == batch_size: | |
| batch_X = torch.cat(batch_list) | |
| outputs = ort_sess.run(None, {'video': batch_X.numpy()}) | |
| y1pred = outputs[0] | |
| y2pred = outputs[1] | |
| y3pred = outputs[2] | |
| y4pred = outputs[3] | |
| for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list): | |
| periodLength = y1.squeeze() | |
| periodicity = y2.squeeze() | |
| marks = y3.squeeze() | |
| event_type = y4.squeeze() | |
| period_lengths[idx:idx+seq_len] += periodLength | |
| periodicities[idx:idx+seq_len] += periodicity | |
| full_marks[idx:idx+seq_len] += marks | |
| event_type_logits[idx:idx+seq_len] += event_type | |
| period_length_overlaps[idx:idx+seq_len] += 1 | |
| event_type_logit_overlaps[idx:idx+seq_len] += 1 | |
| batch_list = [] | |
| idx_list = [] | |
| progress(i / (length + stride_length - stride_pad), desc="Processing...") | |
| 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) | |
| outputs = ort_sess.run(None, {'video': batch_X.numpy()}) | |
| y1pred = outputs[0] | |
| y2pred = outputs[1] | |
| y3pred = outputs[2] | |
| y4pred = outputs[3] | |
| for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list): | |
| periodLength = y1.squeeze() | |
| periodicity = y2.squeeze() | |
| marks = y3.squeeze() | |
| event_type = y4.squeeze() | |
| period_lengths[idx:idx+seq_len] += periodLength | |
| periodicities[idx:idx+seq_len] += periodicity | |
| full_marks[idx:idx+seq_len] += marks | |
| event_type_logits[idx:idx+seq_len] += event_type | |
| period_length_overlaps[idx:idx+seq_len] += 1 | |
| event_type_logit_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] | |
| 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] | |
| 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) | |
| #full_marks_mask = np.int32(full_marks > marks_threshold) | |
| 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) | |
| numofReps = 0 | |
| count = [] | |
| for i in range(len(periodLength)): | |
| if periodLength[i] < 2 or periodicity_mask[i] == 0: | |
| numofReps += 0 | |
| 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 | |
| else: | |
| numofReps += max(0, periodicity_mask[i]/(periodLength[i])) | |
| 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 | |
| confidence = (np.mean(periodicity[periodicity > miss_threshold]) - miss_threshold) / (1 - miss_threshold) | |
| self_err = abs(count_pred - marks_count_pred) | |
| self_pct_err = self_err / count_pred | |
| total_confidence = confidence * (1 - self_pct_err) | |
| if both_feet: | |
| count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}" | |
| else: | |
| count_msg = f"## Predicted Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}" | |
| if api_call: | |
| if count_only_api: | |
| return f"{count_pred:.2f} (conf: {total_confidence:.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(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \ | |
| f"reps: {count_pred:.2f}, marks: {marks_count_pred:.1f}, confidence: {total_confidence:.2f}", \ | |
| f"single_rope_speed: {event_type_probs[0]:.3f}, double_dutch: {event_type_probs[1]:.3f}, double_unders: {event_type_probs[2]:.3f}, single_bounce: {event_type_probs[3]:.3f}" | |
| jumps_per_second = np.clip(1 / ((periodLength / fps) + 0.01), 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, | |
| '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") | |
| 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' | |
| )) | |
| 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]) | |
| return count_msg, fig, hist, bar | |
| 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(theme='WeixuanYuan/Soft_dark') as demo: | |
| gr.Markdown(DESCRIPTION) | |
| in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', | |
| width=400, height=400, interactive=True, container=True, | |
| max_length=150) | |
| 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.Column(min_width=480): | |
| #out_video = gr.PlayableVideo(label="Output Video", elem_id='output-video', format='mp4') | |
| with gr.Row(): | |
| out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot') | |
| 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') | |
| 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!" | |
| 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." | |
| 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 to get the traditional double under count." | |
| faq += "\n* **Q:** Does the model count both feet?\n * **A:** Yes, it counts every time the rope goes around no matter the event." | |
| gr.Markdown(faq) | |
| demo_inference = partial(inference, count_only_api=False, api_key=None) | |
| gr.Examples(examples=[ | |
| [os.path.join(os.path.dirname(__file__), "files", "dylan.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train14.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_17.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train13.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_213.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_156.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_202.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_57.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_95.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_253.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_66.mp4")], | |
| #[os.path.join(os.path.dirname(__file__), "files", "train_21.mp4")] | |
| ], | |
| inputs=[in_video], | |
| outputs=[out_text, out_plot, out_hist, out_event_type_dist], | |
| fn=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces') | |
| run_button.click(demo_inference, [in_video], outputs=[out_text, out_plot, out_hist, out_event_type_dist]) | |
| 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) |