Update to use new outputs

app.py

@@ -10,7 +10,6 @@ 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
@@ -39,87 +38,22 @@ compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU", config=c
 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)
-        # fig, axs = plt.subplots(1, 1, figsize = (3, 3))
 
-        # img_ax = axs
-        # img_ax.imshow(np.zeros((128, 128, 3)))
+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 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 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.85, 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, 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())
@@ -145,44 +79,46 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
             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))
+        all_frames.append(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)
+    output_layer_marks = compiled_model_ir.output(2)
+    output_layer_event_type = compiled_model_ir.output(3)
     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)
+    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))
     for _ in range(seq_len + stride_length):  # pad full sequence
         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(SquarePad())
-                transforms_list.append(transforms.CenterCrop((img_size, img_size)))
-            else:
-                transforms_list.append(transforms.Resize((img_size, img_size)))
+            # 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)))
             
 
             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)
+            frameTensor = preprocess(img).unsqueeze(0)
             Xlist.append(frameTensor)
 
         if len(Xlist) < seq_len:
@@ -198,12 +134,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]
-            for y1, y2, idx in zip(y1pred, y2pred, idx_list):
+            y4pred = result[output_layer_event_type]
+            for y1, y2, y4, idx in zip(y1pred, y2pred, y4pred, idx_list):
                 periodLength = y1.squeeze()
                 periodicity = y2.squeeze()
+                event_type = y4.squeeze()
                 period_lengths[idx:idx+seq_len] += periodLength
                 periodicities[idx:idx+seq_len] += periodicity
+                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 = []
     if len(batch_list) != 0:  # still some leftover frames
@@ -214,15 +154,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]
-        for y1, y2, idx in zip(y1pred, y2pred, idx_list):
+        y4pred = result[output_layer_event_type]
+        for y1, y2, y4, idx in zip(y1pred, y2pred, y4pred, idx_list):
             periodLength = y1.squeeze()
             periodicity = y2.squeeze()
+            event_type = y4.squeeze()
             period_lengths[idx:idx+seq_len] += periodLength
             periodicities[idx:idx+seq_len] += periodicity
+            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]
+    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  
+    event_type_probs = np.exp(event_type_logits) / np.sum(np.exp(event_type_logits))
     
     if median_pred_filter:
         periodicity = medfilt(periodicity, 5)
@@ -252,7 +200,9 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
             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(periodicity, formatter={'float_kind':lambda x: "%.2f" % x}).replace('\n', ''), \
+                f"{count_pred:.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.05), 0, 8)
@@ -305,8 +255,17 @@ def inference(x, count_only_api, api_key, img_size=192, seq_len=64, stride_lengt
                         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])
+    
+    # make a bar plot of the event type distribution
+
+    bar = px.bar(x=['single rope speed', 'double dutch', 'double unders', 'single bounce'], 
+                 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, periodLength
+    return count_msg, fig, hist, bar
         
 
 DESCRIPTION = '# NextJump'
@@ -318,10 +277,10 @@ 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)
+            in_video = gr.Video(label="Input Video", elem_id='input-video', format='mp4', width=400, height=400)
             
     with gr.Row():
-        run_button = gr.Button(label="Run", elem_id='run-button', style=dict(full_width=False), scale=1)
+        run_button = gr.Button(label="Run", elem_id='run-button', 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)
@@ -334,8 +293,13 @@ with gr.Blocks() as demo:
                 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_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
-        out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
+        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:"
@@ -362,10 +326,10 @@ with gr.Blocks() as demo:
                         [b, False, True, -1, True, 1.0, 0.95],
                     ],
                 inputs=[in_video],
-                outputs=[out_text, out_plot, out_hist],
+                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])
+    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')