update to mobilenetv3

app.py

@@ -25,7 +25,7 @@ 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_mobilenetv3.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'])
@@ -46,7 +46,7 @@ else:
     ort_sess = ort.InferenceSession(onnx_file)
 
 print('Warmup...')
-dummy_input = torch.randn(4, 64, 3, 288, 288)   
+dummy_input = torch.randn(4, 64, 3, 224, 224)   
 ort_sess.run(None, {'video': dummy_input.numpy()})
 print('Done!')
 
@@ -73,7 +73,7 @@ def create_transform(img_size):
 
 
 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, 
+              img_size=224, 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()):
@@ -327,48 +327,47 @@ def inference(stream_url, start_time, end_time, count_only_api, api_key,
                  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)
 
-    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__':
+    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():
-                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')
-              
+                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)
 
-    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')
+        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')
 
-if __name__ == '__main__':
     demo.queue(api_open=True, max_size=15).launch(share=False)