add tensorrt if available

app.py

@@ -7,7 +7,6 @@ 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
@@ -26,16 +25,16 @@ 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'])
-# 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"
+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'
 
 # 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)
+# config = {'PERFORMANCE_HINT': 'LATENCY'}
+# compiled_model_ir = ie.compile_model(model=model_ir, device_name='CPU', config=config)
 
 
 class SquarePad:
@@ -59,32 +58,18 @@ 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, 
               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...")
+    progress(0, desc='Starting...')
     x = download_clips(stream_url, os.getcwd(), start_time, end_time)
     # 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)})]
+        providers = ['TensorrtExecutionProvider', ('CUDAExecutionProvider', {'device_id': torch.cuda.current_device(),
+                                                'user_compute_stream': str(torch.cuda.current_stream().cuda_stream)})]
         sess_options = ort.SessionOptions()
         sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
         ort_sess = ort.InferenceSession(onnx_file, sess_options=sess_options, providers=providers)
@@ -96,7 +81,7 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
     if api_call:
         has_access = pbkdf2_sha256.verify(os.environ['DEV_API_TOKEN'], api_key)
         if not has_access:
-            return "Invalid API Key"
+            return 'Invalid API Key'
         
     cap = cv2.VideoCapture(x)
     length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -127,37 +112,26 @@ 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)):
-        X = batch_queue.get()
+        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
         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()})
@@ -178,10 +152,7 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
                 event_type_logit_overlaps[idx:idx+seq_len] += 1
             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()
+        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])
@@ -258,19 +229,19 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
     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}"
+        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}"
+        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})"
+            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}"
+            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)
@@ -303,25 +274,25 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
                     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)",
+                    title='Jumping speed (jumps-per-second)',
                     trendline='rolling',
                     trendline_options=dict(window=16),
-                    trendline_color_override="goldenrod",
+                    trendline_color_override='goldenrod',
                     trendline_scope='overall',
-                    template="plotly_dark")
+                    template='plotly_dark')
     
     fig.update_layout(legend=dict(
-            orientation="h",
-            yanchor="bottom",
+            orientation='h',
+            yanchor='bottom',
             y=0.98,
-            xanchor="right",
+            xanchor='right',
             x=1,
             font=dict(
-                family="Courier",
+                family='Courier',
                 size=12,
-                color="black"
+                color='black'
                 ),
-            bgcolor="AliceBlue",
+            bgcolor='AliceBlue',
         ),
         paper_bgcolor='rgba(0,0,0,0)',
         plot_bgcolor='rgba(0,0,0,0)'
@@ -335,71 +306,64 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
     ))
 
     hist = px.histogram(df, 
-                        x="jumps per second", 
-                        template="plotly_dark", 
-                        marginal="box",
+                        x='jumps per second', 
+                        template='plotly_dark', 
+                        marginal='box',
                         histnorm='percent',
-                        title="Distribution of jumping speed (jumps-per-second)")
+                        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",
+                 template='plotly_dark',
+                 title='Event Type Distribution',
                  labels={'x': 'event type', 'y': 'probability'},
                  range_y=[0, 1])
 
     return x, count_msg, fig, hist, bar
         
 
+with gr.Blocks() as demo:
+    # 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)
 
-
-
-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.Column(elem_id='output-video-container'):
         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)
-                
+                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():
-            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():
+                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.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')
-                
+    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_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')
+
+if __name__ == '__main__':
     demo.queue(api_open=True, max_size=15).launch(share=False)