dylanplummer commited on
Commit
2e41093
·
verified ·
1 Parent(s): 9f1ce44

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +381 -381
app.py CHANGED
@@ -1,382 +1,382 @@
1
- import gradio as gr
2
- import numpy as np
3
- from PIL import Image
4
- from openvino.runtime import Core
5
- import os
6
- import cv2
7
- import uuid
8
- import time
9
- import subprocess
10
- import matplotlib
11
- matplotlib.use('Agg')
12
- import matplotlib.pyplot as plt
13
- from scipy.signal import medfilt
14
- from functools import partial
15
- from passlib.hash import pbkdf2_sha256
16
- from tqdm import tqdm
17
- import pandas as pd
18
- import plotly.express as px
19
- import torch
20
- from torchvision import transforms
21
- import torchvision.transforms.functional as F
22
-
23
- from huggingface_hub import hf_hub_download
24
- from huggingface_hub import HfApi
25
-
26
- plt.style.use('dark_background')
27
-
28
- hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.bin", repo_type="model", token=os.environ['DATASET_SECRET'])
29
- model_xml = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.xml", repo_type="model", token=os.environ['DATASET_SECRET'])
30
- hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.mapping", repo_type="model", token=os.environ['DATASET_SECRET'])
31
- #model_xml = "model_ir/model.xml"
32
-
33
- ie = Core()
34
- model_ir = ie.read_model(model=model_xml)
35
- config = {"PERFORMANCE_HINT": "LATENCY"}
36
- compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)
37
-
38
-
39
- class SquarePad:
40
- # https://discuss.pytorch.org/t/how-to-resize-and-pad-in-a-torchvision-transforms-compose/71850/9
41
- def __call__(self, image):
42
- w, h = image.size
43
- max_wh = max(w, h)
44
- hp = int((max_wh - w) / 2)
45
- vp = int((max_wh - h) / 2)
46
- padding = (hp, vp, hp, vp)
47
- return F.pad(image, padding, 0, 'constant')
48
-
49
- def sigmoid(x):
50
- return 1 / (1 + np.exp(-x))
51
-
52
-
53
- 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.8, marks_threshold=0.6, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
54
- print(x)
55
- #api = HfApi(token=os.environ['DATASET_SECRET'])
56
- #out_file = str(uuid.uuid1())
57
- has_access = False
58
- if api_call:
59
- has_access = pbkdf2_sha256.verify(os.environ['DEV_API_TOKEN'], api_key)
60
- if not has_access:
61
- return "Invalid API Key"
62
-
63
- cap = cv2.VideoCapture(x)
64
- length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
65
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
66
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
67
- period_length_overlaps = np.zeros(length + seq_len)
68
- fps = int(cap.get(cv2.CAP_PROP_FPS))
69
- seconds = length / fps
70
- all_frames = []
71
- frame_i = 1
72
- while cap.isOpened():
73
- ret, frame = cap.read()
74
- if ret is False:
75
- frame = all_frames[-1] # padding will be with last frame
76
- break
77
- frame = cv2.cvtColor(np.uint8(frame), cv2.COLOR_BGR2RGB)
78
- img = Image.fromarray(frame)
79
- all_frames.append(img)
80
- frame_i += 1
81
- cap.release()
82
-
83
- # Get output layer
84
- output_layer_period_length = compiled_model_ir.output(0)
85
- output_layer_periodicity = compiled_model_ir.output(1)
86
- output_layer_marks = compiled_model_ir.output(2)
87
- output_layer_event_type = compiled_model_ir.output(3)
88
- length = len(all_frames)
89
- period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
90
- periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
91
- full_marks = np.zeros(len(all_frames) + seq_len + stride_length)
92
- event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 6))
93
- period_length_overlaps = np.zeros(len(all_frames) + seq_len + stride_length)
94
- event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 6))
95
- for _ in range(seq_len + stride_length): # pad full sequence
96
- all_frames.append(all_frames[-1])
97
- batch_list = []
98
- idx_list = []
99
- for i in tqdm(range(0, length + stride_length - stride_pad, stride_length)):
100
- batch = all_frames[i:i + seq_len]
101
- Xlist = []
102
- for img in batch:
103
- transforms_list = []
104
- # if center_crop:
105
- # if width > height:
106
- # transforms_list.append(transforms.Resize((int(width / (height / img_size)), img_size)))
107
- # else:
108
- # transforms_list.append(transforms.Resize((img_size, int(height / (width / img_size)))))
109
- # transforms_list.append(transforms.CenterCrop((img_size, img_size)))
110
- # else:
111
- transforms_list.append(SquarePad())
112
- transforms_list.append(transforms.Resize((img_size, img_size)))
113
-
114
-
115
- transforms_list += [
116
- transforms.ToTensor()]
117
- #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
118
- preprocess = transforms.Compose(transforms_list)
119
- frameTensor = preprocess(img).unsqueeze(0)
120
- Xlist.append(frameTensor)
121
-
122
- if len(Xlist) < seq_len:
123
- for _ in range(seq_len - len(Xlist)):
124
- Xlist.append(Xlist[-1])
125
-
126
- X = torch.cat(Xlist)
127
- X *= 255
128
- batch_list.append(X.unsqueeze(0))
129
- idx_list.append(i)
130
- if len(batch_list) == batch_size:
131
- batch_X = torch.cat(batch_list)
132
- result = compiled_model_ir(batch_X)
133
- y1pred = result[output_layer_period_length]
134
- y2pred = result[output_layer_periodicity]
135
- y3pred = result[output_layer_marks]
136
- y4pred = result[output_layer_event_type]
137
- for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
138
- periodLength = y1.squeeze()
139
- periodicity = y2.squeeze()
140
- marks = y3.squeeze()
141
- event_type = y4.squeeze()
142
- period_lengths[idx:idx+seq_len] += periodLength
143
- periodicities[idx:idx+seq_len] += periodicity
144
- full_marks[idx:idx+seq_len] += marks
145
- event_type_logits[idx:idx+seq_len] += event_type
146
- period_length_overlaps[idx:idx+seq_len] += 1
147
- event_type_logit_overlaps[idx:idx+seq_len] += 1
148
- batch_list = []
149
- idx_list = []
150
- if len(batch_list) != 0: # still some leftover frames
151
- while len(batch_list) != batch_size:
152
- batch_list.append(batch_list[-1])
153
- idx_list.append(idx_list[-1])
154
- batch_X = torch.cat(batch_list)
155
- result = compiled_model_ir(batch_X)
156
- y1pred = result[output_layer_period_length]
157
- y2pred = result[output_layer_periodicity]
158
- y3pred = result[output_layer_marks]
159
- y4pred = result[output_layer_event_type]
160
- for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
161
- periodLength = y1.squeeze()
162
- periodicity = y2.squeeze()
163
- marks = y3.squeeze()
164
- event_type = y4.squeeze()
165
- period_lengths[idx:idx+seq_len] += periodLength
166
- periodicities[idx:idx+seq_len] += periodicity
167
- full_marks[idx:idx+seq_len] += marks
168
- event_type_logits[idx:idx+seq_len] += event_type
169
- period_length_overlaps[idx:idx+seq_len] += 1
170
- event_type_logit_overlaps[idx:idx+seq_len] += 1
171
-
172
- periodLength = np.divide(period_lengths, period_length_overlaps, where=period_length_overlaps!=0)[:length]
173
- periodicity = np.divide(periodicities, period_length_overlaps, where=period_length_overlaps!=0)[:length]
174
- full_marks = np.divide(full_marks, period_length_overlaps, where=period_length_overlaps!=0)[:length]
175
- per_frame_event_type_logits = np.divide(event_type_logits, event_type_logit_overlaps, where=event_type_logit_overlaps!=0)[:length]
176
- event_type_logits = np.mean(per_frame_event_type_logits, axis=0)
177
- # softmax of event type logits
178
- event_type_probs = np.exp(event_type_logits) / np.sum(np.exp(event_type_logits))
179
- per_frame_event_types = np.argmax(per_frame_event_type_logits, axis=1)
180
-
181
- if median_pred_filter:
182
- periodicity = medfilt(periodicity, 5)
183
- periodLength = medfilt(periodLength, 5)
184
- periodicity = sigmoid(periodicity)
185
- full_marks = sigmoid(full_marks)
186
- full_marks_mask = np.int32(full_marks > marks_threshold)
187
- periodicity_mask = np.int32(periodicity > miss_threshold)
188
- numofReps = 0
189
- count = []
190
- for i in range(len(periodLength)):
191
- if periodLength[i] < 2 or periodicity_mask[i] == 0:
192
- numofReps += 0
193
- else:
194
- numofReps += max(0, periodicity_mask[i]/(periodLength[i]))
195
- count.append(round(float(numofReps), 2))
196
- count_pred = count[-1]
197
- marks_count_pred = 0
198
- for i in range(len(full_marks) - 1):
199
- # if a jump was counted, and periodicity is high, and the next frame was not counted (to avoid double counting)
200
- if full_marks_mask[i] > 0 and periodicity_mask[i] > 0 and full_marks_mask[i + 1] == 0:
201
- marks_count_pred += 1
202
- if not both_feet:
203
- count_pred = count_pred / 2
204
- marks_count_pred = marks_count_pred / 2
205
- count = np.array(count) / 2
206
-
207
- confidence = (np.mean(periodicity[periodicity > miss_threshold]) - miss_threshold) / (1 - miss_threshold)
208
- self_err = abs(count_pred - marks_count_pred)
209
- self_pct_err = self_err / count_pred
210
- total_confidence = confidence * (1 - self_pct_err)
211
-
212
- if both_feet:
213
- count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
214
- else:
215
- count_msg = f"## Predicted Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
216
-
217
- if api_call:
218
- if count_only_api:
219
- return f"{count_pred:.2f} (conf: {total_confidence:.2f})"
220
- else:
221
- return np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
222
- np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
223
- np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
224
- f"reps: {count_pred:.2f}, marks: {marks_count_pred:.1f}, confidence: {total_confidence:.2f}", \
225
- 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}"
226
-
227
-
228
- jumps_per_second = np.clip(1 / ((periodLength / fps) + 0.01), 0, 10)
229
- jumping_speed = np.copy(jumps_per_second)
230
- misses = periodicity < miss_threshold
231
- jumps_per_second[misses] = 0
232
- frame_type = np.array(['miss' if miss else 'frame' for miss in misses])
233
- frame_type[full_marks > marks_threshold] = 'jump'
234
- per_frame_event_types = np.clip(per_frame_event_types, 0, 6) / 6
235
- df = pd.DataFrame.from_dict({'period length': periodLength,
236
- 'jumping speed': jumping_speed,
237
- 'jumps per second': jumps_per_second,
238
- 'periodicity': periodicity,
239
- 'miss': misses,
240
- 'frame_type': frame_type,
241
- 'event_type': per_frame_event_types,
242
- 'jumps': full_marks,
243
- 'jumps_size': (full_marks + 0.05) * 10,
244
- 'miss_size': np.clip((1 - periodicity) * 0.9 + 0.1, 1, 10),
245
- 'seconds': np.linspace(0, seconds, num=len(periodLength))})
246
- event_type_tick_vals = np.linspace(0, 1, num=7)
247
- event_type_colors = ['red', 'orange', 'green', 'blue', 'purple', 'pink', 'black']
248
- fig = px.scatter(data_frame=df,
249
- x='seconds',
250
- y='jumps per second',
251
- #symbol='frame_type',
252
- #symbol_map={'frame': 'circle', 'miss': 'circle-open', 'jump': 'triangle-down'},
253
- color='event_type',
254
- size='jumps_size',
255
- size_max=10,
256
- color_continuous_scale=[(t, c) for t, c in zip(event_type_tick_vals, event_type_colors)],
257
- range_color=(0,1),
258
- title="Jumping speed (jumps-per-second)",
259
- trendline='rolling',
260
- trendline_options=dict(window=16),
261
- trendline_color_override="goldenrod",
262
- trendline_scope='overall',
263
- template="plotly_dark")
264
-
265
- fig.update_layout(legend=dict(
266
- orientation="h",
267
- yanchor="bottom",
268
- y=0.98,
269
- xanchor="right",
270
- x=1,
271
- font=dict(
272
- family="Courier",
273
- size=12,
274
- color="black"
275
- ),
276
- bgcolor="AliceBlue",
277
- ),
278
- paper_bgcolor='rgba(0,0,0,0)',
279
- plot_bgcolor='rgba(0,0,0,0)'
280
- )
281
- fig.update_layout(coloraxis_colorbar=dict(
282
- tickvals=event_type_tick_vals,
283
- ticktext=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'],
284
- title='event type'
285
- ))
286
-
287
- hist = px.histogram(df,
288
- x="jumps per second",
289
- template="plotly_dark",
290
- marginal="box",
291
- histnorm='percent',
292
- title="Distribution of jumping speed (jumps-per-second)")
293
-
294
- # make a bar plot of the event type distribution
295
-
296
- bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders'],
297
- y=event_type_probs,
298
- template="plotly_dark",
299
- title="Event Type Distribution",
300
- labels={'x': 'event type', 'y': 'probability'},
301
- range_y=[0, 1])
302
-
303
- return count_msg, fig, hist, bar
304
-
305
-
306
- DESCRIPTION = '# NextJump 🦘'
307
- DESCRIPTION += '\n## AI Counting for Competitive Jump Rope'
308
- 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).'
309
-
310
-
311
- with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
312
- gr.Markdown(DESCRIPTION)
313
- with gr.Column():
314
- #with gr.Row():
315
- in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', width='50%', min_width=400, interactive=True, container=True)
316
- placeholder = gr.Markdown(label="", elem_id='placeholder-text')
317
-
318
- with gr.Row():
319
- run_button = gr.Button(value="Run", elem_id='run-button', scale=1)
320
- api_dummy_button = gr.Button(value="Run (No Viz)", elem_id='count-only', visible=False, scale=2)
321
- count_only = gr.Checkbox(label="Count Only", visible=False)
322
- api_token = gr.Textbox(label="API Key", elem_id='api-token', visible=False)
323
-
324
- with gr.Column(elem_id='output-video-container'):
325
- with gr.Row():
326
- with gr.Column():
327
- out_text = gr.Markdown(label="Predicted Count", elem_id='output-text')
328
- period_length = gr.Textbox(label="Period Length", elem_id='period-length', visible=False)
329
- periodicity = gr.Textbox(label="Periodicity", elem_id='periodicity', visible=False)
330
- #with gr.Column(min_width=480):
331
- #out_video = gr.PlayableVideo(label="Output Video", elem_id='output-video', format='mp4')
332
- with gr.Row():
333
- out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
334
- with gr.Row():
335
- with gr.Column():
336
- out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
337
- with gr.Column():
338
- out_event_type_dist = gr.Plot(label="Event Type Distribution", elem_id='output-event-type-dist')
339
-
340
- with gr.Accordion(label="Instructions and more information", open=False):
341
- instructions = "## Instructions:"
342
- 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!"
343
- instructions += "\n\n## Tips (optional):"
344
- instructions += "\n* Trim the video to start and end of the event"
345
- instructions += "\n* Frame the jumper fully, in the center of the frame"
346
- 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."
347
- gr.Markdown(instructions)
348
-
349
- faq = "## FAQ:"
350
- 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."
351
- 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."
352
- 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."
353
- faq += "\n* **Q:** Does the model count both feet?\n * **A:** Yes, it counts every time the rope goes around no matter the event."
354
- gr.Markdown(faq)
355
-
356
- demo_inference = partial(inference, count_only_api=False, api_key=None)
357
-
358
- gr.Examples(examples=[
359
- [os.path.join(os.path.dirname(__file__), "files", "dylan.mp4")],
360
- [os.path.join(os.path.dirname(__file__), "files", "train14.mp4")],
361
- [os.path.join(os.path.dirname(__file__), "files", "train_17.mp4")],
362
- #[os.path.join(os.path.dirname(__file__), "files", "train13.mp4")],
363
- #[os.path.join(os.path.dirname(__file__), "files", "train_213.mp4")],
364
- [os.path.join(os.path.dirname(__file__), "files", "train_156.mp4")],
365
- #[os.path.join(os.path.dirname(__file__), "files", "train_202.mp4")],
366
- [os.path.join(os.path.dirname(__file__), "files", "train_57.mp4")],
367
- [os.path.join(os.path.dirname(__file__), "files", "train_95.mp4")],
368
- [os.path.join(os.path.dirname(__file__), "files", "train_253.mp4")],
369
- [os.path.join(os.path.dirname(__file__), "files", "train_66.mp4")],
370
- [os.path.join(os.path.dirname(__file__), "files", "train_21.mp4")]
371
- ],
372
- inputs=[in_video],
373
- outputs=[out_text, out_plot, out_hist, out_event_type_dist],
374
- fn=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
375
-
376
- run_button.click(demo_inference, [in_video], outputs=[out_text, out_plot, out_hist, out_event_type_dist])
377
- api_inference = partial(inference, api_call=True)
378
- api_dummy_button.click(api_inference, [in_video, count_only, api_token], outputs=[period_length], api_name='inference')
379
-
380
-
381
- if __name__ == "__main__":
382
  demo.queue(api_open=True, max_size=15).launch(share=False)
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ from PIL import Image
4
+ from openvino.runtime import Core
5
+ import os
6
+ import cv2
7
+ import uuid
8
+ import time
9
+ import subprocess
10
+ import matplotlib
11
+ matplotlib.use('Agg')
12
+ import matplotlib.pyplot as plt
13
+ from scipy.signal import medfilt
14
+ from functools import partial
15
+ from passlib.hash import pbkdf2_sha256
16
+ from tqdm import tqdm
17
+ import pandas as pd
18
+ import plotly.express as px
19
+ import torch
20
+ from torchvision import transforms
21
+ import torchvision.transforms.functional as F
22
+
23
+ from huggingface_hub import hf_hub_download
24
+ from huggingface_hub import HfApi
25
+
26
+ plt.style.use('dark_background')
27
+
28
+ hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.bin", repo_type="model", token=os.environ['DATASET_SECRET'])
29
+ model_xml = hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.xml", repo_type="model", token=os.environ['DATASET_SECRET'])
30
+ hf_hub_download(repo_id="dylanplummer/ropenet", filename="model.mapping", repo_type="model", token=os.environ['DATASET_SECRET'])
31
+ #model_xml = "model_ir/model.xml"
32
+
33
+ ie = Core()
34
+ model_ir = ie.read_model(model=model_xml)
35
+ config = {"PERFORMANCE_HINT": "LATENCY"}
36
+ compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=config)
37
+
38
+
39
+ class SquarePad:
40
+ # https://discuss.pytorch.org/t/how-to-resize-and-pad-in-a-torchvision-transforms-compose/71850/9
41
+ def __call__(self, image):
42
+ w, h = image.size
43
+ max_wh = max(w, h)
44
+ hp = int((max_wh - w) / 2)
45
+ vp = int((max_wh - h) / 2)
46
+ padding = (hp, vp, hp, vp)
47
+ return F.pad(image, padding, 0, 'constant')
48
+
49
+ def sigmoid(x):
50
+ return 1 / (1 + np.exp(-x))
51
+
52
+
53
+ 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.8, marks_threshold=0.6, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
54
+ print(x)
55
+ #api = HfApi(token=os.environ['DATASET_SECRET'])
56
+ #out_file = str(uuid.uuid1())
57
+ has_access = False
58
+ if api_call:
59
+ has_access = pbkdf2_sha256.verify(os.environ['DEV_API_TOKEN'], api_key)
60
+ if not has_access:
61
+ return "Invalid API Key"
62
+
63
+ cap = cv2.VideoCapture(x)
64
+ length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
65
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
66
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
67
+ period_length_overlaps = np.zeros(length + seq_len)
68
+ fps = int(cap.get(cv2.CAP_PROP_FPS))
69
+ seconds = length / fps
70
+ all_frames = []
71
+ frame_i = 1
72
+ while cap.isOpened():
73
+ ret, frame = cap.read()
74
+ if ret is False:
75
+ frame = all_frames[-1] # padding will be with last frame
76
+ break
77
+ frame = cv2.cvtColor(np.uint8(frame), cv2.COLOR_BGR2RGB)
78
+ img = Image.fromarray(frame)
79
+ all_frames.append(img)
80
+ frame_i += 1
81
+ cap.release()
82
+
83
+ # Get output layer
84
+ output_layer_period_length = compiled_model_ir.output(0)
85
+ output_layer_periodicity = compiled_model_ir.output(1)
86
+ output_layer_marks = compiled_model_ir.output(2)
87
+ output_layer_event_type = compiled_model_ir.output(3)
88
+ length = len(all_frames)
89
+ period_lengths = np.zeros(len(all_frames) + seq_len + stride_length)
90
+ periodicities = np.zeros(len(all_frames) + seq_len + stride_length)
91
+ full_marks = np.zeros(len(all_frames) + seq_len + stride_length)
92
+ event_type_logits = np.zeros((len(all_frames) + seq_len + stride_length, 7))
93
+ period_length_overlaps = np.zeros(len(all_frames) + seq_len + stride_length)
94
+ event_type_logit_overlaps = np.zeros((len(all_frames) + seq_len + stride_length, 7))
95
+ for _ in range(seq_len + stride_length): # pad full sequence
96
+ all_frames.append(all_frames[-1])
97
+ batch_list = []
98
+ idx_list = []
99
+ for i in tqdm(range(0, length + stride_length - stride_pad, stride_length)):
100
+ batch = all_frames[i:i + seq_len]
101
+ Xlist = []
102
+ for img in batch:
103
+ transforms_list = []
104
+ # if center_crop:
105
+ # if width > height:
106
+ # transforms_list.append(transforms.Resize((int(width / (height / img_size)), img_size)))
107
+ # else:
108
+ # transforms_list.append(transforms.Resize((img_size, int(height / (width / img_size)))))
109
+ # transforms_list.append(transforms.CenterCrop((img_size, img_size)))
110
+ # else:
111
+ transforms_list.append(SquarePad())
112
+ transforms_list.append(transforms.Resize((img_size, img_size)))
113
+
114
+
115
+ transforms_list += [
116
+ transforms.ToTensor()]
117
+ #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
118
+ preprocess = transforms.Compose(transforms_list)
119
+ frameTensor = preprocess(img).unsqueeze(0)
120
+ Xlist.append(frameTensor)
121
+
122
+ if len(Xlist) < seq_len:
123
+ for _ in range(seq_len - len(Xlist)):
124
+ Xlist.append(Xlist[-1])
125
+
126
+ X = torch.cat(Xlist)
127
+ X *= 255
128
+ batch_list.append(X.unsqueeze(0))
129
+ idx_list.append(i)
130
+ if len(batch_list) == batch_size:
131
+ batch_X = torch.cat(batch_list)
132
+ result = compiled_model_ir(batch_X)
133
+ y1pred = result[output_layer_period_length]
134
+ y2pred = result[output_layer_periodicity]
135
+ y3pred = result[output_layer_marks]
136
+ y4pred = result[output_layer_event_type]
137
+ for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
138
+ periodLength = y1.squeeze()
139
+ periodicity = y2.squeeze()
140
+ marks = y3.squeeze()
141
+ event_type = y4.squeeze()
142
+ period_lengths[idx:idx+seq_len] += periodLength
143
+ periodicities[idx:idx+seq_len] += periodicity
144
+ full_marks[idx:idx+seq_len] += marks
145
+ event_type_logits[idx:idx+seq_len] += event_type
146
+ period_length_overlaps[idx:idx+seq_len] += 1
147
+ event_type_logit_overlaps[idx:idx+seq_len] += 1
148
+ batch_list = []
149
+ idx_list = []
150
+ if len(batch_list) != 0: # still some leftover frames
151
+ while len(batch_list) != batch_size:
152
+ batch_list.append(batch_list[-1])
153
+ idx_list.append(idx_list[-1])
154
+ batch_X = torch.cat(batch_list)
155
+ result = compiled_model_ir(batch_X)
156
+ y1pred = result[output_layer_period_length]
157
+ y2pred = result[output_layer_periodicity]
158
+ y3pred = result[output_layer_marks]
159
+ y4pred = result[output_layer_event_type]
160
+ for y1, y2, y3, y4, idx in zip(y1pred, y2pred, y3pred, y4pred, idx_list):
161
+ periodLength = y1.squeeze()
162
+ periodicity = y2.squeeze()
163
+ marks = y3.squeeze()
164
+ event_type = y4.squeeze()
165
+ period_lengths[idx:idx+seq_len] += periodLength
166
+ periodicities[idx:idx+seq_len] += periodicity
167
+ full_marks[idx:idx+seq_len] += marks
168
+ event_type_logits[idx:idx+seq_len] += event_type
169
+ period_length_overlaps[idx:idx+seq_len] += 1
170
+ event_type_logit_overlaps[idx:idx+seq_len] += 1
171
+
172
+ periodLength = np.divide(period_lengths, period_length_overlaps, where=period_length_overlaps!=0)[:length]
173
+ periodicity = np.divide(periodicities, period_length_overlaps, where=period_length_overlaps!=0)[:length]
174
+ full_marks = np.divide(full_marks, period_length_overlaps, where=period_length_overlaps!=0)[:length]
175
+ per_frame_event_type_logits = np.divide(event_type_logits, event_type_logit_overlaps, where=event_type_logit_overlaps!=0)[:length]
176
+ event_type_logits = np.mean(per_frame_event_type_logits, axis=0)
177
+ # softmax of event type logits
178
+ event_type_probs = np.exp(event_type_logits) / np.sum(np.exp(event_type_logits))
179
+ per_frame_event_types = np.argmax(per_frame_event_type_logits, axis=1)
180
+
181
+ if median_pred_filter:
182
+ periodicity = medfilt(periodicity, 5)
183
+ periodLength = medfilt(periodLength, 5)
184
+ periodicity = sigmoid(periodicity)
185
+ full_marks = sigmoid(full_marks)
186
+ full_marks_mask = np.int32(full_marks > marks_threshold)
187
+ periodicity_mask = np.int32(periodicity > miss_threshold)
188
+ numofReps = 0
189
+ count = []
190
+ for i in range(len(periodLength)):
191
+ if periodLength[i] < 2 or periodicity_mask[i] == 0:
192
+ numofReps += 0
193
+ else:
194
+ numofReps += max(0, periodicity_mask[i]/(periodLength[i]))
195
+ count.append(round(float(numofReps), 2))
196
+ count_pred = count[-1]
197
+ marks_count_pred = 0
198
+ for i in range(len(full_marks) - 1):
199
+ # if a jump was counted, and periodicity is high, and the next frame was not counted (to avoid double counting)
200
+ if full_marks_mask[i] > 0 and periodicity_mask[i] > 0 and full_marks_mask[i + 1] == 0:
201
+ marks_count_pred += 1
202
+ if not both_feet:
203
+ count_pred = count_pred / 2
204
+ marks_count_pred = marks_count_pred / 2
205
+ count = np.array(count) / 2
206
+
207
+ confidence = (np.mean(periodicity[periodicity > miss_threshold]) - miss_threshold) / (1 - miss_threshold)
208
+ self_err = abs(count_pred - marks_count_pred)
209
+ self_pct_err = self_err / count_pred
210
+ total_confidence = confidence * (1 - self_pct_err)
211
+
212
+ if both_feet:
213
+ count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
214
+ else:
215
+ count_msg = f"## Predicted Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}, Confidence: {total_confidence:.2f}"
216
+
217
+ if api_call:
218
+ if count_only_api:
219
+ return f"{count_pred:.2f} (conf: {total_confidence:.2f})"
220
+ else:
221
+ return np.array2string(periodLength, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
222
+ np.array2string(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
223
+ np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
224
+ f"reps: {count_pred:.2f}, marks: {marks_count_pred:.1f}, confidence: {total_confidence:.2f}", \
225
+ 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}"
226
+
227
+
228
+ jumps_per_second = np.clip(1 / ((periodLength / fps) + 0.01), 0, 10)
229
+ jumping_speed = np.copy(jumps_per_second)
230
+ misses = periodicity < miss_threshold
231
+ jumps_per_second[misses] = 0
232
+ frame_type = np.array(['miss' if miss else 'frame' for miss in misses])
233
+ frame_type[full_marks > marks_threshold] = 'jump'
234
+ per_frame_event_types = np.clip(per_frame_event_types, 0, 7) / 7
235
+ df = pd.DataFrame.from_dict({'period length': periodLength,
236
+ 'jumping speed': jumping_speed,
237
+ 'jumps per second': jumps_per_second,
238
+ 'periodicity': periodicity,
239
+ 'miss': misses,
240
+ 'frame_type': frame_type,
241
+ 'event_type': per_frame_event_types,
242
+ 'jumps': full_marks,
243
+ 'jumps_size': (full_marks + 0.05) * 10,
244
+ 'miss_size': np.clip((1 - periodicity) * 0.9 + 0.1, 1, 10),
245
+ 'seconds': np.linspace(0, seconds, num=len(periodLength))})
246
+ event_type_tick_vals = np.linspace(0, 1, num=7)
247
+ event_type_colors = ['red', 'orange', 'green', 'blue', 'purple', 'pink', 'black']
248
+ fig = px.scatter(data_frame=df,
249
+ x='seconds',
250
+ y='jumps per second',
251
+ #symbol='frame_type',
252
+ #symbol_map={'frame': 'circle', 'miss': 'circle-open', 'jump': 'triangle-down'},
253
+ color='event_type',
254
+ size='jumps_size',
255
+ size_max=10,
256
+ color_continuous_scale=[(t, c) for t, c in zip(event_type_tick_vals, event_type_colors)],
257
+ range_color=(0,1),
258
+ title="Jumping speed (jumps-per-second)",
259
+ trendline='rolling',
260
+ trendline_options=dict(window=16),
261
+ trendline_color_override="goldenrod",
262
+ trendline_scope='overall',
263
+ template="plotly_dark")
264
+
265
+ fig.update_layout(legend=dict(
266
+ orientation="h",
267
+ yanchor="bottom",
268
+ y=0.98,
269
+ xanchor="right",
270
+ x=1,
271
+ font=dict(
272
+ family="Courier",
273
+ size=12,
274
+ color="black"
275
+ ),
276
+ bgcolor="AliceBlue",
277
+ ),
278
+ paper_bgcolor='rgba(0,0,0,0)',
279
+ plot_bgcolor='rgba(0,0,0,0)'
280
+ )
281
+ fig.update_layout(coloraxis_colorbar=dict(
282
+ tickvals=event_type_tick_vals,
283
+ ticktext=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'],
284
+ title='event type'
285
+ ))
286
+
287
+ hist = px.histogram(df,
288
+ x="jumps per second",
289
+ template="plotly_dark",
290
+ marginal="box",
291
+ histnorm='percent',
292
+ title="Distribution of jumping speed (jumps-per-second)")
293
+
294
+ # make a bar plot of the event type distribution
295
+
296
+ bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounces', 'double bounces', 'triple unders', 'other'],
297
+ y=event_type_probs,
298
+ template="plotly_dark",
299
+ title="Event Type Distribution",
300
+ labels={'x': 'event type', 'y': 'probability'},
301
+ range_y=[0, 1])
302
+
303
+ return count_msg, fig, hist, bar
304
+
305
+
306
+ DESCRIPTION = '# NextJump 🦘'
307
+ DESCRIPTION += '\n## AI Counting for Competitive Jump Rope'
308
+ 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).'
309
+
310
+
311
+ with gr.Blocks(theme='WeixuanYuan/Soft_dark') as demo:
312
+ gr.Markdown(DESCRIPTION)
313
+ with gr.Column():
314
+ #with gr.Row():
315
+ in_video = gr.PlayableVideo(label="Input Video", elem_id='input-video', format='mp4', width='50%', min_width=400, interactive=True, container=True)
316
+ placeholder = gr.Markdown(label="", elem_id='placeholder-text')
317
+
318
+ with gr.Row():
319
+ run_button = gr.Button(value="Run", elem_id='run-button', scale=1)
320
+ api_dummy_button = gr.Button(value="Run (No Viz)", elem_id='count-only', visible=False, scale=2)
321
+ count_only = gr.Checkbox(label="Count Only", visible=False)
322
+ api_token = gr.Textbox(label="API Key", elem_id='api-token', visible=False)
323
+
324
+ with gr.Column(elem_id='output-video-container'):
325
+ with gr.Row():
326
+ with gr.Column():
327
+ out_text = gr.Markdown(label="Predicted Count", elem_id='output-text')
328
+ period_length = gr.Textbox(label="Period Length", elem_id='period-length', visible=False)
329
+ periodicity = gr.Textbox(label="Periodicity", elem_id='periodicity', visible=False)
330
+ #with gr.Column(min_width=480):
331
+ #out_video = gr.PlayableVideo(label="Output Video", elem_id='output-video', format='mp4')
332
+ with gr.Row():
333
+ out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
334
+ with gr.Row():
335
+ with gr.Column():
336
+ out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
337
+ with gr.Column():
338
+ out_event_type_dist = gr.Plot(label="Event Type Distribution", elem_id='output-event-type-dist')
339
+
340
+ with gr.Accordion(label="Instructions and more information", open=False):
341
+ instructions = "## Instructions:"
342
+ 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!"
343
+ instructions += "\n\n## Tips (optional):"
344
+ instructions += "\n* Trim the video to start and end of the event"
345
+ instructions += "\n* Frame the jumper fully, in the center of the frame"
346
+ 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."
347
+ gr.Markdown(instructions)
348
+
349
+ faq = "## FAQ:"
350
+ 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."
351
+ 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."
352
+ 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."
353
+ faq += "\n* **Q:** Does the model count both feet?\n * **A:** Yes, it counts every time the rope goes around no matter the event."
354
+ gr.Markdown(faq)
355
+
356
+ demo_inference = partial(inference, count_only_api=False, api_key=None)
357
+
358
+ gr.Examples(examples=[
359
+ [os.path.join(os.path.dirname(__file__), "files", "dylan.mp4")],
360
+ [os.path.join(os.path.dirname(__file__), "files", "train14.mp4")],
361
+ [os.path.join(os.path.dirname(__file__), "files", "train_17.mp4")],
362
+ #[os.path.join(os.path.dirname(__file__), "files", "train13.mp4")],
363
+ #[os.path.join(os.path.dirname(__file__), "files", "train_213.mp4")],
364
+ [os.path.join(os.path.dirname(__file__), "files", "train_156.mp4")],
365
+ #[os.path.join(os.path.dirname(__file__), "files", "train_202.mp4")],
366
+ [os.path.join(os.path.dirname(__file__), "files", "train_57.mp4")],
367
+ [os.path.join(os.path.dirname(__file__), "files", "train_95.mp4")],
368
+ [os.path.join(os.path.dirname(__file__), "files", "train_253.mp4")],
369
+ [os.path.join(os.path.dirname(__file__), "files", "train_66.mp4")],
370
+ [os.path.join(os.path.dirname(__file__), "files", "train_21.mp4")]
371
+ ],
372
+ inputs=[in_video],
373
+ outputs=[out_text, out_plot, out_hist, out_event_type_dist],
374
+ fn=demo_inference, cache_examples=os.getenv('SYSTEM') == 'spaces')
375
+
376
+ run_button.click(demo_inference, [in_video], outputs=[out_text, out_plot, out_hist, out_event_type_dist])
377
+ api_inference = partial(inference, api_call=True)
378
+ api_dummy_button.click(api_inference, [in_video, count_only, api_token], outputs=[period_length], api_name='inference')
379
+
380
+
381
+ if __name__ == "__main__":
382
  demo.queue(api_open=True, max_size=15).launch(share=False)