multiprocessing

app.py

@@ -7,6 +7,7 @@ import math
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
+import multiprocessing as mp
 from scipy.signal import medfilt, find_peaks
 from functools import partial
 from passlib.hash import pbkdf2_sha256
@@ -25,8 +26,8 @@ from hls_download import download_clips
 
 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'])
-onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="nextjump_fp16.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="nextjump.onnx", repo_type="model", token=os.environ['DATASET_SECRET'])
+#onnx_file = hf_hub_download(repo_id="dylanplummer/ropenet", filename="nextjump_fp16.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"
@@ -58,6 +59,20 @@ def create_transform(img_size):
         transforms.ToTensor(),
     ])
 
+def preprocess_frame(img, img_size):
+    preprocess = create_transform(img_size)
+    frameTensor = preprocess(img).unsqueeze(0)
+    return frameTensor * 255
+
+
+def worker_function(frame_queue, batch_queue, img_size, seq_len):
+    while True:
+        frames = frame_queue.get()
+        if frames is None:  # Signal to exit
+            break
+        batch = torch.cat([preprocess_frame(img, img_size) for img in frames])
+        batch_queue.put(batch)
+
 
 def inference(stream_url, start_time, end_time, count_only_api, api_key, 
               img_size=288, seq_len=64, stride_length=32, stride_pad=3, batch_size=4, 
@@ -112,29 +127,40 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
     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])
+
+    num_workers = mp.cpu_count()  # Use all available CPU cores
+    frame_queue = mp.Queue(maxsize=num_workers * 2)
+    batch_queue = mp.Queue(maxsize=num_workers * 2)
+
+     # Start worker processes
+    processes = []
+    for _ in range(num_workers):
+        p = mp.Process(target=worker_function, args=(frame_queue, batch_queue, img_size, seq_len))
+        p.start()
+        processes.append(p)
+
+    # Enqueue frame batches
+    for i in tqdm(range(0, length + stride_length - stride_pad, stride_length)):
+        batch = all_frames[i:i + seq_len]
+        if len(batch) < seq_len:
+            batch.extend([batch[-1]] * (seq_len - len(batch)))
+        frame_queue.put(batch)
+
+    # Signal workers to exit after all frames are processed
+    for _ in range(num_workers):
+        frame_queue.put(None)
+
     batch_list = []
     idx_list = []
-    preprocess = create_transform(img_size)
+    #preprocess = create_transform(img_size)
     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:
-            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
+        X = batch_queue.get()
         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': np.float16(batch_X.numpy())})
+            outputs = ort_sess.run(None, {'video': batch_X.numpy()})
             y1pred = outputs[0]
             y2pred = outputs[1]
             y3pred = outputs[2]
@@ -153,12 +179,15 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
             batch_list = []
             idx_list = []
         progress(i / (length + stride_length - stride_pad), desc="Processing...")
+    # Wait for all processes to finish
+    for p in processes:
+        p.join()
     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': np.float16(batch_X.numpy())})
+        outputs = ort_sess.run(None, {'video': batch_X.numpy()})
         y1pred = outputs[0]
         y2pred = outputs[1]
         y3pred = outputs[2]
@@ -324,52 +353,53 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
     return x, 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():
-        in_stream_url = gr.Textbox(label="Stream URL", elem_id='stream-url', visible=True)
-        with gr.Column():
-            in_stream_start = gr.Textbox(label="Start Time", elem_id='stream-start', visible=True)
-        with gr.Column():
-            in_stream_end = gr.Textbox(label="End Time", elem_id='stream-end', visible=True)
-        with gr.Column(min_width=480):
-            out_video = gr.PlayableVideo(label="Video Clip", elem_id='output-video', format='mp4', width=400, height=400)
-            
-    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.Row():
-            out_plot = gr.Plot(label="Jumping Speed", elem_id='output-plot')
+
+
+if __name__ == "__main__":
+    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)
+        # 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():
+            in_stream_url = gr.Textbox(label="Stream URL", elem_id='stream-url', visible=True)
             with gr.Column():
-                out_hist = gr.Plot(label="Speed Histogram", elem_id='output-hist')
+                in_stream_start = gr.Textbox(label="Start Time", elem_id='stream-start', visible=True)
             with gr.Column():
-                out_event_type_dist = gr.Plot(label="Event Type Distribution", elem_id='output-event-type-dist')
-              
-
-    demo_inference = partial(inference, count_only_api=False, api_key=None)
-    
-    run_button.click(demo_inference, [in_stream_url, in_stream_start, in_stream_end], outputs=[out_video, out_text, out_plot, out_hist, out_event_type_dist])
-    api_inference = partial(inference, api_call=True)
-    api_dummy_button.click(api_inference, [in_stream_url, in_stream_start, in_stream_end, count_only, api_token], outputs=[period_length], api_name='inference')
+                in_stream_end = gr.Textbox(label="End Time", elem_id='stream-end', visible=True)
+            with gr.Column(min_width=480):
+                out_video = gr.PlayableVideo(label="Video Clip", elem_id='output-video', format='mp4', width=400, height=400)
+                
+        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.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')
+                
 
-if __name__ == "__main__":
+        demo_inference = partial(inference, count_only_api=False, api_key=None)
+        
+        run_button.click(demo_inference, [in_stream_url, in_stream_start, in_stream_end], outputs=[out_video, out_text, out_plot, out_hist, out_event_type_dist])
+        api_inference = partial(inference, api_call=True)
+        api_dummy_button.click(api_inference, [in_stream_url, in_stream_start, in_stream_end, count_only, api_token], outputs=[period_length], api_name='inference')
     demo.queue(api_open=True, max_size=15).launch(share=False)