update to use marks

app.py

@@ -11,6 +11,7 @@ import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 from scipy.signal import medfilt
+from skimage.measure import block_reduce
 from functools import partial
 from passlib.hash import pbkdf2_sha256
 from tqdm import tqdm
@@ -28,7 +29,7 @@ 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"
+model_xml = "model_ir/model.xml"
 
 ie = Core()
 model_ir = ie.read_model(model=model_xml)
@@ -53,7 +54,7 @@ def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 
 
-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, median_pred_filter=True, center_crop=True, both_feet=True, api_call=False):
+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):
     print(x)
     #api = HfApi(token=os.environ['DATASET_SECRET'])
     #out_file = str(uuid.uuid1())
@@ -91,6 +92,7 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     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, 4))
     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, 4))
@@ -134,13 +136,16 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
             result = compiled_model_ir(batch_X)
             y1pred = result[output_layer_period_length]
             y2pred = result[output_layer_periodicity]
+            y3pred = result[output_layer_marks]
             y4pred = result[output_layer_event_type]
-            for y1, y2, y4, idx in zip(y1pred, y2pred, y4pred, idx_list):
+            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
@@ -154,19 +159,23 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
         result = compiled_model_ir(batch_X)
         y1pred = result[output_layer_period_length]
         y2pred = result[output_layer_periodicity]
+        y3pred = result[output_layer_marks]
         y4pred = result[output_layer_event_type]
-        for y1, y2, y4, idx in zip(y1pred, y2pred, y4pred, idx_list):
+        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]
     event_type_logits = np.divide(event_type_logits, event_type_logit_overlaps, where=event_type_logit_overlaps!=0)[:length]
     event_type_logits = np.mean(event_type_logits, axis=0)
     # softmax of event type logits  
@@ -174,9 +183,13 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     
     if median_pred_filter:
         periodicity = medfilt(periodicity, 5)
-        periodLength = medfilt(periodLength, 5)
+        #periodLength = medfilt(periodLength, 5)
     periodicity = sigmoid(periodicity)
+    full_marks = sigmoid(full_marks)
+    full_marks_mask = np.int32(full_marks > marks_threshold)
+    #full_marks_reduced = block_reduce(full_marks > marks_threshold, (3,), np.max)
     periodicity_mask = np.int32(periodicity > miss_threshold)
+    #periodicity_mask_reduced = block_reduce(periodicity_mask, (3,), np.max)
     numofReps = 0
     count = []
     for i in range(len(periodLength)):
@@ -186,14 +199,20 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
             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
 
     if both_feet:
-        count_msg = f"## Predicted Count (both feet): {count_pred:.1f}"
+        count_msg = f"## Reps Count (both feet): {count_pred:.1f}, Marks Count (both feet): {marks_count_pred:.1f}"
     else:
-        count_msg = f"## Predicted Count (one foot): {count_pred:.1f}"
+        count_msg = f"## Predicted Count (one foot): {count_pred:.1f}, Marks Count (one foot): {marks_count_pred:.1f}"
 
     if api_call:
         if count_only_api:
@@ -201,11 +220,12 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
         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', ''), \
-                f"{count_pred:.2f}", \
+                np.array2string(full_marks, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
+                f"reps: {count_pred:.2f}, marks: {marks_count_pred:.1f}", \
                 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.05), 0, 8)
+    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
@@ -214,20 +234,22 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
                                 'jumps per second': jumps_per_second,
                                 'periodicity': periodicity,
                                 'miss': misses,
+                                'jumps': full_marks,
+                                'jumps_size': (full_marks + 0.2) * 10,
                                 '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'},
+                    symbol_map={False: 'circle', True: 'circle-open'},
                     color='periodicity',
-                    size='miss_size',
-                    size_max=8,
+                    size='jumps_size',
+                    size_max=10,
                     color_continuous_scale='RdYlGn',
                     title="Jumping speed (jumps-per-second)",
                     trendline='rolling',
-                    trendline_options=dict(window=32),
+                    trendline_options=dict(window=16),
                     trendline_color_override="goldenrod",
                     trendline_scope='overall',
                     template="plotly_dark")
@@ -267,7 +289,7 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
     return count_msg, fig, hist, bar
         
 
-DESCRIPTION = '# NextJump'
+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).'
 

requirements.txt

@@ -4,6 +4,7 @@ matplotlib
 plotly
 passlib
 scipy
+scikit-image
 --find-links https://download.pytorch.org/whl/torch_stable.html
 opencv-python-headless==4.7.0.68
 openvino-dev==2022.3.0