Update app.py

app.py

@@ -15,13 +15,12 @@ from denseav.plotting import plot_attention_video, plot_2head_attention_video, p
 from denseav.shared import norm, crop_to_divisor, blur_dim
 from os.path import join
 
-
 if __name__ == "__main__":
 
-    os.environ['TORCH_HOME'] = '/tmp/.cache'
-    os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
-    sample_images_dir = "/tmp/samples"
-    # sample_videos_dir = "samples"
+    # os.environ['TORCH_HOME'] = '/tmp/.cache'
+    # os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
+    # sample_images_dir = "/tmp/samples"
+    sample_videos_dir = "samples"
 
 
     def download_video(url, save_path):
@@ -33,6 +32,10 @@ if __name__ == "__main__":
     base_url = "https://marhamilresearch4.blob.core.windows.net/denseav-public/samples/"
     sample_videos_urls = {
         "puppies.mp4": base_url + "puppies.mp4",
+        "peppers.mp4": base_url + "peppers.mp4",
+        "boat.mp4": base_url + "boat.mp4",
+        "elephant2.mp4": base_url + "elephant2.mp4",
+
     }
 
     # Ensure the directory for sample videos exists
@@ -49,7 +52,7 @@ if __name__ == "__main__":
             print(f"{filename} already exists. Skipping download.")
 
     csv.field_size_limit(100000000)
-    options = ['language', "sound", "sound_and_language"]
+    options = ['language', "sound_and_language", "sound"]
     load_size = 224
     plot_size = 224
 
@@ -145,21 +148,41 @@ if __name__ == "__main__":
         )
         return temp_video_path_1, temp_video_path_2, temp_video_path_3, temp_video_path_4
 
+        return temp_video_path_1, temp_video_path_2, temp_video_path_3
+
 
     with gr.Blocks() as demo:
         with gr.Column():
-            video_input.render()
-            model_option.render()
+            gr.Markdown("## Visualizing Sound and Language with DenseAV")
+            gr.Markdown(
+                "This demo allows you to explore the inner attention maps of DenseAV's dense multi-head contrastive operator.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    model_option.render()
+                with gr.Column(scale=3):
+                    video_input.render()
+            with gr.Row():
+                submit_button = gr.Button("Submit")
+            with gr.Row():
+                gr.Examples(
+                    examples=[
+                        [join(sample_videos_dir, "puppies.mp4"), "sound_and_language"],
+                        [join(sample_videos_dir, "peppers.mp4"), "language"],
+                        [join(sample_videos_dir, "elephant2.mp4"), "language"],
+                        [join(sample_videos_dir, "boat.mp4"), "language"]
+
+                    ],
+                    inputs=[video_input, model_option]
+                )
             with gr.Row():
                 video_output1.render()
                 video_output2.render()
-            with gr.Row():
                 video_output3.render()
-                video_output4.render()
 
-        demo.examples = [
-            [join(sample_videos_dir, "puppies.mp4"), "language"],
-        ]
+        submit_button.click(fn=process_video, inputs=[video_input, model_option],
+                            outputs=[video_output1, video_output2])
 
     # demo.launch(server_name="0.0.0.0", server_port=6006, debug=True)
-    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)
+
+    demo.launch(server_name="0.0.0.0", server_port=6006, debug=True)
+    # demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)